Upload 225 files

.gitattributes

@@ -33,3 +33,23 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+llava/model/coat/optimizer/kernels/build/lib.linux-x86_64-cpython-310/qoptim_cuda.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+llava/model/coat/optimizer/kernels/build/temp.linux-x86_64-cpython-310/bindings.o filter=lfs diff=lfs merge=lfs -text
+llava/model/coat/optimizer/kernels/build/temp.linux-x86_64-cpython-310/fp8_adamw_cuda.o filter=lfs diff=lfs merge=lfs -text
+llava/model/coat/optimizer/kernels/build/temp.linux-x86_64-cpython-310/fp8_adamw_expand_cuda.o filter=lfs diff=lfs merge=lfs -text
+static/af3_main_diagram-1.png filter=lfs diff=lfs merge=lfs -text
+static/af3_radial-1.png filter=lfs diff=lfs merge=lfs -text
+static/af3_sota.png filter=lfs diff=lfs merge=lfs -text
+static/audio/audio2.wav filter=lfs diff=lfs merge=lfs -text
+static/chat/audio1.mp3 filter=lfs diff=lfs merge=lfs -text
+static/chat/audio2.mp3 filter=lfs diff=lfs merge=lfs -text
+static/emergent/audio1.wav filter=lfs diff=lfs merge=lfs -text
+static/logo-no-bg.png filter=lfs diff=lfs merge=lfs -text
+static/speech/339a1acd-afcb-466b-a7b1-8661e59b1e56.wav filter=lfs diff=lfs merge=lfs -text
+static/speech/audio3.wav filter=lfs diff=lfs merge=lfs -text
+static/speech/bcc6057d-0dda-435d-b956-a96ab27bc9e4.wav filter=lfs diff=lfs merge=lfs -text
+static/speech/be84d293-5e9c-4158-9a1e-b4dd1acb7d70.wav filter=lfs diff=lfs merge=lfs -text
+static/speech/fec3402e-7883-45c0-90d4-38647f615dc3.wav filter=lfs diff=lfs merge=lfs -text
+static/think/audio1.wav filter=lfs diff=lfs merge=lfs -text
+static/think/audio2.wav filter=lfs diff=lfs merge=lfs -text
+static/voice/voice_2.mp3 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,14 +1,139 @@
----
-title: Audio Flamingo 3
-emoji: ⚡
-colorFrom: pink
-colorTo: purple
-sdk: gradio
-sdk_version: 5.36.2
-app_file: app.py
-pinned: false
-license: other
-short_description: Online demo for Audio Flamingo 3
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+<div align="center" style="display: flex; justify-content: center; align-items: center; text-align: center;">
+  <a href="https://github.com/NVIDIA/audio-flamingo" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
+    <img src="static/logo-no-bg.png" alt="Audio Flamingo 3 🔥🚀🔥" width="120">
+  </a>
+</div>
+<div align="center" style="display: flex; justify-content: center; align-items: center; text-align: center;">
+    <h2>
+    Audio Flamingo 3: Advancing Audio Intelligence with Fully Open Large Audio-Language Models
+    </h2>
+</div>
+
+<div align="center" style="display: flex; justify-content: center; margin-top: 10px;">
+  <a href=""><img src="https://img.shields.io/badge/arXiv-2503.03983-AD1C18" style="margin-right: 5px;"></a>
+  <a href="https://research.nvidia.com/labs/adlr/AF3/"><img src="https://img.shields.io/badge/Demo page-228B22" style="margin-right: 5px;"></a>
+  <a href="https://github.com/NVIDIA/audio-flamingo"><img src='https://img.shields.io/badge/Github-Audio Flamingo 3-9C276A' style="margin-right: 5px;"></a>
+  <a href="https://github.com/NVIDIA/audio-flamingo/stargazers"><img src="https://img.shields.io/github/stars/NVIDIA/audio-flamingo.svg?style=social"></a>
+</div>
+
+<div align="center" style="display: flex; justify-content: center; margin-top: 10px; flex-wrap: wrap; gap: 5px;">
+  <a href="https://huggingface.co/nvidia/audio-flamingo-3">
+    <img src="https://img.shields.io/badge/🤗-Checkpoints-ED5A22.svg">
+  </a>
+  <a href="https://huggingface.co/nvidia/audio-flamingo-3-chat">
+    <img src="https://img.shields.io/badge/🤗-Checkpoints (Chat)-ED5A22.svg">
+  </a>
+  <a href="https://huggingface.co/datasets/nvidia/AudioSkills">
+    <img src="https://img.shields.io/badge/🤗-Dataset: AudioSkills--XL-ED5A22.svg">
+  </a>
+  <a href="https://huggingface.co/datasets/nvidia/LongAudio">
+    <img src="https://img.shields.io/badge/🤗-Dataset: LongAudio--XL-ED5A22.svg">
+  </a>
+  <a href="https://huggingface.co/datasets/nvidia/AF-Chat">
+    <img src="https://img.shields.io/badge/🤗-Dataset: AF--Chat-ED5A22.svg">
+  </a>
+  <a href="https://huggingface.co/datasets/nvidia/AF-Think">
+    <img src="https://img.shields.io/badge/🤗-Dataset: AF--Think-ED5A22.svg">
+  </a>
+</div>
+
+<div align="center" style="display: flex; justify-content: center; margin-top: 10px;">
+<a href="https://huggingface.co/spaces/nvidia/audio_flamingo_3"><img src="https://img.shields.io/badge/🤗-Gradio Demo (7B)-5F9EA0.svg" style="margin-right: 5px;"></a>
+</div>
+
+## Overview
+
+This repo contains the PyTorch implementation of [Audio Flamingo 3: Advancing Audio Intelligence with Fully Open Large Audio-Language Models](). Audio Flamingo 3 (AF3) is a fully open, state-of-the-art Large Audio-Language Model (LALM) that advances reasoning and understanding across speech, sounds, and music. AF3 builds on previous work with innovations in:
+
+- Unified audio representation learning (speech, sound, music)  
+- Flexible, on-demand chain-of-thought reasoning (Thinking in Audio) 
+- Long-context audio comprehension (including speech and up to 10 minutes)
+- Multi-turn, multi-audio conversational dialogue (AF3-Chat)    
+- Voice-to-voice interaction (AF3-Chat)    
+
+Extensive evaluations confirm AF3’s effectiveness, setting new benchmarks on over 20 public audio understanding and reasoning tasks.
+
+
+## Main Results
+
+Audio Flamingo 3 outperforms prior SOTA models including GAMA, Audio Flamingo, Audio Flamingo 2, Qwen-Audio, Qwen2-Audio, Qwen2.5-Omni.LTU, LTU-AS, SALMONN, AudioGPT, Gemini Flash v2 and Gemini Pro v1.5 on a number of understanding and reasoning benchmarks.
+
+<div align="center">
+  <img class="img-full" src="static/af3_radial-1.png" width="300">
+</div>
+
+<div align="center">
+  <img class="img-full" src="static/af3_sota.png" width="400">
+</div>
+
+## Audio Flamingo 3 Architecture
+
+Audio Flamingo 3 uses AF-Whisper unified audio encoder, MLP-based audio adaptor, Decoder-only LLM backbone (Qwen2.5-7B), and Streaming TTS module (AF3-Chat).
+Audio Flamingo 3 can take up to 10 minutes of audio inputs. 
+
+<div align="center">
+  <img class="img-full" src="static/af3_main_diagram-1.png" width="800">
+</div>
+
+## Installation
+
+```bash
+./environment_setup.sh af3
+```
+
+## Code Structure
+
+- The folder ```audio_flamingo_3/``` contains the main training and inference code of Audio Flamingo 3.
+- The folder ```audio_flamingo_3/scripts``` contains the inference scripts of Audio Flamingo 3 in case you would like to use our pretrained checkpoints on HuggingFace.
+
+Each folder is self-contained and we expect no cross dependencies between these folders. This repo does not contain the code for Streaming-TTS pipeline which will released in the near future.
+
+## Single Line Inference
+
+To infer stage 3 model directly, run the command below:
+```bash
+python llava/cli/infer_audio.py --model-base /path/to/checkpoint/af3-7b --conv-mode auto --text "Please describe the audio in detail" --media static/audio1.wav
+```
+
+To infer the model in stage 3.5 model, run the command below:
+```bash
+python llava/cli/infer_audio.py --model-base /path/to/checkpoint/af3-7b --model-path /path/to/checkpoint/af3-7b/stage35 --conv-mode auto --text "Please describe the audio in detail" --media static/audio1.wav --peft-mode
+```
+
+## References
+
+The main training and inferencing code within each folder are modified from [NVILA](https://github.com/NVlabs/VILA/tree/main) [Apache license](incl_licenses/License_1.md).
+
+## License
+
+- The code in this repo is under [MIT license](incl_licenses/MIT_license.md).
+- The checkpoints are for non-commercial use only [NVIDIA OneWay Noncommercial License](incl_licenses/NVIDIA_OneWay_Noncommercial_License.docx). They are also subject to the [Qwen Research license](https://huggingface.co/Qwen/Qwen2.5-7B/blob/main/LICENSE), the [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and the original licenses accompanying each training dataset.
+- Notice: Audio Flamingo 3 is built with Qwen-2.5. Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) Alibaba Cloud. All Rights Reserved.
+
+
+## Citation
+
+- Audio Flamingo 2
+```
+@article{ghosh2025audio,
+  title={Audio Flamingo 2: An Audio-Language Model with Long-Audio Understanding and Expert Reasoning Abilities},
+  author={Ghosh, Sreyan and Kong, Zhifeng and Kumar, Sonal and Sakshi, S and Kim, Jaehyeon and Ping, Wei and Valle, Rafael and Manocha, Dinesh and Catanzaro, Bryan},
+  journal={arXiv preprint arXiv:2503.03983},
+  year={2025}
+}
+```
+
+- Audio Flamingo
+```
+@inproceedings{kong2024audio,
+  title={Audio Flamingo: A Novel Audio Language Model with Few-Shot Learning and Dialogue Abilities},
+  author={Kong, Zhifeng and Goel, Arushi and Badlani, Rohan and Ping, Wei and Valle, Rafael and Catanzaro, Bryan},
+  booktitle={International Conference on Machine Learning},
+  pages={25125--25148},
+  year={2024},
+  organization={PMLR}
+}
+```
+
+

llava/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+from .entry import *
+from .media import *

llava/cli/infer_audio.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import argparse
+import importlib.util
+import json
+import os
+
+from pydantic import BaseModel
+from termcolor import colored
+
+import llava
+from llava import conversation as clib
+from llava.media import Image, Video, Sound
+from llava.model.configuration_llava import JsonSchemaResponseFormat, ResponseFormat
+from peft import PeftModel
+import torch
+
+def get_schema_from_python_path(path: str) -> str:
+    schema_path = os.path.abspath(path)
+    spec = importlib.util.spec_from_file_location("schema_module", schema_path)
+    schema_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(schema_module)
+
+    # Get the Main class from the loaded module
+    Main = schema_module.Main
+    assert issubclass(
+        Main, BaseModel
+    ), f"The provided python file {path} does not contain a class Main that describes a JSON schema"
+    return Main.schema_json()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-base", "-mb", type=str, required=True)
+    parser.add_argument("--model-path", "-mp", type=str, required=True)
+    parser.add_argument("--conv-mode", "-c", type=str, default="auto")
+    parser.add_argument("--text", type=str)
+    parser.add_argument("--media", type=str, nargs="+")
+    parser.add_argument("--json-mode", action="store_true")
+    parser.add_argument("--peft-mode", action="store_true")
+    parser.add_argument("--json-schema", type=str, default=None)
+    args = parser.parse_args()
+
+    # Convert json mode to response format
+    if not args.json_mode:
+        response_format = None
+    elif args.json_schema is None:
+        response_format = ResponseFormat(type="json_object")
+    else:
+        schema_str = get_schema_from_python_path(args.json_schema)
+        print(schema_str)
+        response_format = ResponseFormat(type="json_schema", json_schema=JsonSchemaResponseFormat(schema=schema_str))
+
+    # Load model
+    model = llava.load(args.model_base)
+    if args.peft_mode:
+        model = PeftModel.from_pretrained(
+            model,
+            args.model_path,
+            device_map="auto",
+            torch_dtype=torch.float16,
+        )
+    # Set conversation mode
+    clib.default_conversation = clib.conv_templates[args.conv_mode].copy()
+
+    # Prepare multi-modal prompt
+    prompt = []
+    if args.media is not None:
+        for media in args.media or []:
+            if any(media.endswith(ext) for ext in [".wav",".mp3", ".flac"]):
+                media = Sound(media)
+            else:
+                raise ValueError(f"Unsupported media type: {media}")
+            prompt.append(media)
+    if args.text is not None:
+        prompt.append(args.text)
+
+    # Generate response
+    response = model.generate_content(prompt, response_format=response_format)
+    print(colored(response, "cyan", attrs=["bold"]))
+
+
+if __name__ == "__main__":
+    main()

llava/constants.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# This file is modified from https://github.com/haotian-liu/LLaVA/
+
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+DEFAULT_SOUND_TOKEN = "<sound>"
+DEFAULT_SPEECH_TOKEN = "<speech>"
+SENTINEL_TOKEN = "<vila/sentinel>"
+
+MEDIA_TOKENS = {
+    "speech": "<speech>",
+    "sound": "<sound>",
+}
+
+
+"""
+151643: AddedToken("<|endoftext|>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151644: AddedToken("<|im_start|>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151645: AddedToken("<|im_end|>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151646: AddedToken("[BOS]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151647: AddedToken("[PAD]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151648: AddedToken("<vila/sentinel>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151649: AddedToken("<image>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151650: AddedToken("<vila/video>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151651: AddedToken("<sound>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151652: AddedToken("<speech>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+
+"""
+NUM_EXTRA_TOKENS = 10

llava/conversation.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+# This file is modified from https://github.com/haotian-liu/LLaVA/
+
+import dataclasses
+from enum import Enum, auto
+from typing import List
+
+from llava.utils.logging import logger
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+
+    AUTO = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_3 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    sep_style: SeparatorStyle = SeparatorStyle.AUTO
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.LLAMA_3:
+            ret = self.system + self.sep
+            for rid, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message = message[0]
+                    sep = self.sep if rid < len(messages) - 1 else self.sep2
+                    ret += role + message + sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version,
+        )
+
+
+conv_auto = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(),
+    sep_style=SeparatorStyle.AUTO,
+    sep="\n",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(),
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+hermes_2 = Conversation(
+    system="<|im_start|>system\nAnswer the questions.",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+    messages=(),
+    version="hermes-2",
+)
+
+# Template added by Yukang. Note (kentang-mit@): sep is <|eot_id|> for official template.
+llama_3_chat = Conversation(
+    system="<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful language and vision assistant. "
+    "You are able to understand the visual content that the user provides, "
+    "and assist the user with a variety of tasks using natural language.",
+    roles=("<|start_header_id|>user<|end_header_id|>\n\n", "<|start_header_id|>assistant<|end_header_id|>\n\n"),
+    version="llama_v3",
+    messages=(),
+    sep_style=SeparatorStyle.LLAMA_3,
+    sep="<|eot_id|>",
+    sep2="<|end_of_text|>",
+)
+
+
+default_conversation = conv_auto
+conv_templates = {
+    "auto": conv_auto,
+    "hermes-2": hermes_2,
+    "llama_3": llama_3_chat,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "plain": conv_llava_plain,
+}
+
+
+CONVERSATION_MODE_MAPPING = {
+    "vila1.5-3b": "vicuna_v1",
+    "vila1.5-8b": "llama_3",
+    "vila1.5-13b": "vicuna_v1",
+    "vila1.5-40b": "hermes-2",
+    "llama-3": "llama_3",
+    "llama3": "llama_3",
+}
+
+
+def auto_set_conversation_mode(model_name_or_path: str) -> str:
+    global default_conversation
+    for k, v in CONVERSATION_MODE_MAPPING.items():
+        if k in model_name_or_path.lower():
+            logger.info(f"Setting conversation mode to `{v}` based on model name/path `{model_name_or_path}`.")
+            default_conversation = conv_templates[v]
+            return

llava/data/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+from .builder import *
+from .dataset import *
+from .datasets_mixture import *

llava/data/base.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import random
+from typing import Any, Dict, List
+
+import torch
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer
+
+from llava.mm_utils import dynamic_process_images_and_prompt, dynamic_s2_process_images_and_prompt, process_images
+from llava.train.args import DataArguments
+from llava.utils.logging import logger
+from llava.utils.media import extract_media
+from llava.utils.tokenizer import preprocess_conversation
+
+__all__ = ["BaseDataset"]
+
+def _process_speech(speech: List[Any], data_args: DataArguments) -> torch.Tensor:
+    return torch.tensor(speech)
+
+def _process_sound(sound: List[Any], data_args: DataArguments) -> torch.Tensor:
+    return torch.tensor(sound)
+
+def _process_sound_masks(sound_masks: List[Any], data_args: DataArguments) -> torch.Tensor:
+    return torch.tensor(sound_masks)
+
+
+class BaseDataset(Dataset):
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizer,
+        data_args: DataArguments,
+        no_system_prompt: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+        self.no_system_prompt = no_system_prompt
+        self.instances = []
+        self.enable_dynamic_res = False
+        self.enable_dynamic_res_s2 = False
+        # global_batch_size: int,
+        self.global_batch_size = kwargs.get("global_batch_size", 1)
+
+        # by default, dataset cls will resample on failure
+        self.resample_on_failure = kwargs.get("resample_on_failure", True)
+
+        # by default, dataset cls will resample on failure
+        self.resample_on_failure = kwargs.get("resample_on_failure", True)
+
+    def process(self, instance: Dict[str, Any]) -> List[Dict[str, Any]]:
+        raise NotImplementedError
+
+    def __getitem__(self, index: int) -> Dict[str, Any]:
+        instance = self.instances[index]
+
+        try:
+            # Process instance to conversation
+            conversation = self.process(instance)
+
+            # Extract media from conversation
+            media, media_meta = extract_media(conversation, self.data_args)
+
+            if "speech" in media:
+                processed_speech = _process_speech(media["speech"], self.data_args)
+            if "sound" in media:
+                processed_sound = _process_sound(media["sound"], self.data_args)
+                processed_sound_feature_masks = _process_sound_masks(media_meta["sound_feature_masks"], self.data_args)
+                processed_sound_embed_masks = _process_sound_masks(media_meta["sound_embed_masks"], self.data_args)
+            # Prepare "input_ids" and "labels" for training
+            data = preprocess_conversation(conversation, self.tokenizer, no_system_prompt=self.no_system_prompt)
+
+            if "speech" in media:
+                data["speech"] = processed_speech
+            if "sound" in media:
+                data["sound"] = processed_sound
+                data["sound_feature_masks"] = processed_sound_feature_masks
+                data["sound_embed_masks"] = processed_sound_embed_masks
+
+        except Exception as e:
+            if not self.resample_on_failure:
+                raise e
+            else:
+                logger.exception(f"Error processing instance '{instance}': '{e}'. Resampling.")
+                return self.__getitem__(random.randint(0, len(self.instances) - 1))
+
+        return data
+
+    def __len__(self) -> int:
+        return len(self.instances)

llava/data/builder.py

@@ -0,0 +1,193 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import os
+import os.path as osp
+from itertools import chain
+from typing import Any, List, Optional
+
+import torch
+import torch.distributed as dist
+from hydra.utils import instantiate
+from torch.utils.data import ConcatDataset, Dataset
+from transformers import PreTrainedTokenizer
+
+from llava.data.datasets_mixture import DATASETS_LEGACY
+from llava.train.args import DataArguments, TrainingArguments
+from llava.utils import io
+from llava.utils.logging import logger
+import time
+import numpy as np
+__all__ = ["DATASETS", "MIXTURES", "register_datasets", "register_mixtures", "parse_mixture", "build_dataset"]
+
+
+def load_dataset_yaml(name):
+    fname = f"{name}.yaml" if not name.endswith(".yaml") else name
+
+    # yaml under llava/data/registry/datasets
+    repo_path = osp.join(osp.dirname(__file__), "registry", "datasets", fname)
+    if osp.exists(repo_path):
+        return repo_path
+
+    # # yaml under <fs yaml path>
+    abs_path = osp.expanduser(fname)
+    if osp.exists(abs_path):
+        return abs_path
+
+    raise FileNotFoundError(f"Dataset '{name}' is not found in the {repo_path} or {abs_path}.")
+
+
+def register_datasets(name: Optional[str] = None):
+    if name is None:
+        name = os.environ.get("VILA_DATASETS", "default")
+        logger.info(f"Registering datasets from environment: '{name}'.")
+    # return io.load(osp.join(osp.dirname(__file__), "registry", "datasets", f"{name}.yaml"))
+    dataset_meta = {}
+    for _name in name.split(","):
+        yamlpath = load_dataset_yaml(_name)
+        logger.info(f"Registering datasets from: '{yamlpath}'.")
+        meta = io.load(yamlpath)
+        dataset_meta.update(meta)
+    return dataset_meta
+
+
+def register_mixtures():
+    return io.load(os.path.join(os.path.dirname(__file__), "registry", "mixtures.yaml"))
+
+
+DATASETS = register_datasets()
+MIXTURES = register_mixtures()
+
+
+def parse_mixture(mixture: str) -> List[str]:
+    names = mixture.split("+") if "+" in mixture else [mixture]
+    while any(name in MIXTURES for name in names):
+        names = list(chain(*[MIXTURES.get(name, [name]) for name in names]))
+    return sorted(names)
+
+
+class SubsetDataset(Dataset):
+    def __init__(self, dataset: Dataset, limit: int) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.limit = limit
+
+    def __len__(self) -> int:
+        return int(len(self.dataset) * self.limit)
+
+    def __getitem__(self, index: int) -> Any:
+        return self.dataset[index % len(self.dataset)]
+
+class RepeatedDataset(Dataset):
+    def __init__(self, dataset: Dataset, times: int) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.times = times
+
+    def __len__(self) -> int:
+        return len(self.dataset) * self.times
+
+    def __getitem__(self, index: int) -> Any:
+        return self.dataset[index % len(self.dataset)]
+
+
+def get_world_size():
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_world_size()
+    else:
+        return 1
+
+
+def build_dataset(
+    mixture: str,
+    data_args: DataArguments,
+    training_args: TrainingArguments,
+    tokenizer: PreTrainedTokenizer,
+) -> Dataset:
+    logger.warning(f"Training VILA with mixture '{mixture}'.")
+    datasets = []
+    dataset_rng = np.random.default_rng(1234)
+    for name in parse_mixture(mixture):        
+
+        if "*" in name:
+            name, times = name.split("*")
+            times = int(times)
+        else:
+            times = 1
+        limit_dataset = False
+        if "#" in name:
+            # we limit the max length of this dataset
+            name, max_length_percent = name.split("#")
+            limit_dataset = True
+        if DATASETS is not None and name in DATASETS:
+            if name in DATASETS_LEGACY:
+                logger.warning(f"Dataset '{name}' exists in both new and legacy registries. Using the new one.")
+            dataset = instantiate(DATASETS[name], _partial_=True)(
+                tokenizer=tokenizer,
+                data_args=data_args,
+                global_batch_size=(
+                    training_args.per_device_train_batch_size
+                    # * torch.distributed.get_world_size()
+                    * get_world_size()
+                    * training_args.gradient_accumulation_steps
+                ),
+            )
+        elif name in DATASETS_LEGACY:
+            logger.warning(f"Dataset '{name}' is from the legacy registry. Please consider migrating it.")
+            dataset = build_dataset_legacy(
+                name,
+                data_args=data_args,
+                training_args=training_args,
+                tokenizer=tokenizer,
+            )
+        else:
+            raise ValueError(f"Dataset '{name}' is not found in the registries.")
+
+        
+        if limit_dataset:
+            # we limit the max length of this dataset
+            max_length = int(float(int(max_length_percent) / 100.) * len(dataset))
+            dataset = SubsetDataset(dataset, float(int(max_length_percent) / 100.))
+
+        if times > 1:
+            dataset = RepeatedDataset(dataset, times)
+        datasets.append(dataset)
+    return ConcatDataset(datasets)
+
+
+def build_dataset_legacy(
+    name: str,
+    data_args: DataArguments,
+    training_args: TrainingArguments,
+    tokenizer: PreTrainedTokenizer,
+) -> Dataset:
+    from llava.data.dataset import (
+        LazySupervisedDataset,
+        LazyWDSDataset,
+    )
+
+    dataset = DATASETS_LEGACY[name]
+    dataset_type = dataset.dataset_type
+    if dataset_type == "torch":
+        dataset_cls = LazySupervisedDataset
+    elif dataset_type == "wds":
+        dataset_cls = LazyWDSDataset
+    else:
+        raise NotImplementedError(f"{dataset_type} is not supported.")
+
+    data_args.meta_path = getattr(dataset, "meta_path", None)
+    data_args.caption_choice = getattr(dataset, "caption_choice", None)
+    data_args.caption_choice_2 = getattr(dataset, "caption_choice_2", None)
+    data_args.start_idx = getattr(dataset, "start_idx", None)
+    data_args.end_idx = getattr(dataset, "end_idx", None)
+
+    return dataset_cls(
+        tokenizer=tokenizer,
+        data_path=dataset.data_path,
+        image_folder=getattr(dataset, "image_path"),
+        data_args=data_args,
+        training_args=training_args,
+    )

llava/data/collate.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Sequence
+
+import torch
+from transformers import PreTrainedTokenizer
+
+from llava.constants import IGNORE_INDEX
+from llava.utils.logging import logger
+
+__all__ = ["DataCollator"]
+
+
+@dataclass
+class DataCollator:
+    tokenizer: PreTrainedTokenizer
+
+    def __init__(self, tokenizer: PreTrainedTokenizer):
+        super().__init__()
+        self.tokenizer = tokenizer
+
+    def __call__(self, instances: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
+        # Gather everything from the batch
+        input_ids, labels, media, block_sizes = [], [], {name: [] for name in self.tokenizer.media_tokens}, []
+
+        media_meta = {}
+
+        media_meta["sound_feature_masks"] = []
+        media_meta["sound_embed_masks"] = []
+        media_meta["frame_times"] = []
+        for instance in instances:
+            if isinstance(instance["input_ids"], torch.Tensor):
+                input_ids.append(instance["input_ids"])
+                labels.append(instance["labels"])
+                for name in media:
+                    objs = instance.get(name)
+                    objs = objs if objs is not None else []
+                    media[name].append([obj for obj in objs])
+                if instance.get("sound") is not None:
+                    for name_k in media_meta:
+                        if "sound" in name_k:
+                            objs = instance.get(name_k)
+                            media_meta[name_k].append([obj for obj in objs])
+                if instance.get("video") is not None or instance.get("image") is not None:
+                    for name_k in media_meta:
+                        if "frame" in name_k:
+                            objs = instance.get(name_k)
+                            media_meta[name_k].append([obj for obj in objs])
+                if "block_sizes" in instance:
+                    block_sizes.append(instance["block_sizes"])
+                else:
+                    block_sizes.append(
+                        [None for _ in range(len(instance.get("image")))] if instance.get("image") is not None else []
+                    )
+            else:
+                input_ids.extend(instance["input_ids"])
+                labels.extend(instance["labels"])
+                for name in media:
+                    objs = instance.get(name)
+                    objs = objs if objs is not None else [[] for _ in range(len(instance["input_ids"]))]
+                    media[name].extend(objs)
+                if instance.get("sound") is not None:
+                    for name_k in media_meta:
+                        if "sound" in name_k:
+                            objs = instance.get(name_k)
+                            media_meta[name_k].extend(objs)
+                if instance.get("video") is not None or instance.get("image") is not None:
+                    for name_k in media_meta:
+                        if "frame" in name_k:
+                            objs = instance.get(name_k)
+                            media_meta[name_k].append([obj for obj in objs])
+                if "block_sizes" in instance:
+                    block_sizes.extend(instance["block_sizes"])
+                else:
+                    block_sizes.extend(
+                        [[None for _ in range(len(objs))] for objs in instance.get("image")]
+                        if instance.get("image") is not None
+                        else [[] for _ in range(len(instance["input_ids"]))]
+                    )
+
+        batch_size = len(input_ids)
+        
+
+        # Check if the number of media objects (or the number of block sizes) matches the number of media tokens
+        for name in media:
+            for k in range(batch_size):
+                if name == "image" and not all([_ is None for _ in block_sizes[k]]):
+                    actual = len(block_sizes[k])
+                else:
+                    actual = len(media[name][k])
+                expected = (input_ids[k] == self.tokenizer.media_token_ids[name]).sum().item()
+                if actual != expected:
+                    raise ValueError(
+                        f"Number mismatch between {name} objects and {name} tokens. "
+                        f"There are {expected} {name} tokens but {actual} {name} objects."
+                    )
+        
+        # Batchify the inputs
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids,
+            batch_first=True,
+            padding_value=self.tokenizer.pad_token_id,
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels,
+            batch_first=True,
+            padding_value=IGNORE_INDEX,
+        )
+        input_ids = input_ids[:, : self.tokenizer.model_max_length]
+        labels = labels[:, : self.tokenizer.model_max_length]
+        attention_mask = input_ids.ne(self.tokenizer.pad_token_id)
+
+        # Truncate media objects if necessary
+        for name in media:
+            objects = []
+            for k in range(batch_size):
+                if name == "image" and not all([_ is None for _ in block_sizes[k]]):
+                    actual = len(media[name][k])
+                    num_large_scale_blocks = sum([x * y for x, y in block_sizes[k]])
+                    num_small_scale_blocks = actual - num_large_scale_blocks
+                    num_small_scale_blocks_each_img = num_small_scale_blocks // len(block_sizes[k])
+                    expected_full_image = (input_ids[k] == self.tokenizer.media_token_ids[name]).sum().item()
+                    expected = (
+                        sum([x * y for x, y in block_sizes[k][:expected_full_image]])
+                        + num_small_scale_blocks_each_img * expected_full_image
+                    )
+                    if actual > expected:
+                        logger.warning(f"Truncating the number of {name} objects from {actual} to {expected}")
+                        media[name][k] = media[name][k][:expected]
+                    objects.extend(media[name][k])
+                    block_sizes[k] = block_sizes[k][:expected_full_image]
+                else:
+                    actual = len(media[name][k])
+                    expected = (input_ids[k] == self.tokenizer.media_token_ids[name]).sum().item()
+                    if actual > expected:
+                        logger.warning(f"Truncating the number of {name} objects from {actual} to {expected}")
+                        media[name][k] = media[name][k][:expected]
+                    objects.extend(media[name][k])
+                    if name == "image":
+                        block_sizes[k] = block_sizes[k][:expected]
+            media[name] = objects
+
+        for name in media_meta:
+            objects = []
+            for k in range(batch_size):
+                try:
+                    objects.extend(media_meta[name][k])
+                except:
+                    continue
+            media_meta[name] = objects
+     
+        # Flatten block sizes from [[bls_im1_instance1, bls_im2_instance1], [bls_im1_instance2, bls_im2_instance2], ...] to [bls_im1_instance1, bls_im2_instance1, bls_im1_instance2, bls_im2_instance2, ...]
+        block_sizes = sum(block_sizes, [])
+        return {
+            "input_ids": input_ids,
+            "media": media,
+            "media_config": {"image": {"block_sizes": block_sizes}, "video": {}, "speech": {}, "sound": {}},
+            "labels": labels,
+            "attention_mask": attention_mask,
+            "media_meta": media_meta,
+        }

llava/data/dataset.py

@@ -0,0 +1,1635 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import base64
+import copy
+import io
+import json
+import os
+import os.path as osp
+import random
+import time
+import warnings
+from dataclasses import dataclass
+from typing import Dict, Sequence
+import math
+import numpy as np
+import PIL
+import torch
+import transformers
+from PIL import Image, ImageFile
+from torch.utils.data import Dataset, default_collate
+from transformers import PreTrainedTokenizer
+from transformers import AutoFeatureExtractor
+import kaldiio
+import llava.data.datasets_mixture as datasets_mixture
+from llava import conversation as conversation_lib
+from llava.constants import DEFAULT_SOUND_TOKEN,DEFAULT_SPEECH_TOKEN, IGNORE_INDEX
+from llava.data.collate import DataCollator
+from llava.mm_utils import (
+    load_audio,
+    get_num_windows,
+    tokenizer_image_token,
+)
+from torchvision import transforms
+from llava.train.args import DataArguments, TrainingArguments
+from llava.train.sequence_parallel import (
+    extract_local_from_list,
+    extract_local_input_ids,
+    extract_local_position_ids,
+    get_pg_manager,
+)
+from llava.utils.tokenizer import preprocess_conversation
+# import torchaudio
+from pytorchvideo.data.clip_sampling import ConstantClipsPerVideoSampler, UniformClipSampler
+import soundfile as sf
+from librosa import resample as librosa_resample
+import whisper
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+PIL.Image.MAX_IMAGE_PIXELS = 1000000000
+
+def int16_to_float32(x):
+    return (x / 32767.0).astype(np.float32)
+
+
+def float32_to_int16(x):
+    x = np.clip(x, a_min=-1., a_max=1.)
+    return (x * 32767.).astype(np.int16)
+
+
+
+def preprocess_multimodal(sources: Sequence[str], data_args: DataArguments) -> Dict:
+    is_multimodal = data_args.is_multimodal
+    if not is_multimodal:
+        return sources
+
+    for source in sources:
+        concat_values = "".join([sentence["value"] for sentence in source])
+        for sid, sentence in enumerate(source):
+            # In multimodal conversations, we automatically prepend '<image>' at the start of the first sentence if it doesn't already contain one.
+            
+            if DEFAULT_SOUND_TOKEN in sentence["value"]:
+                sentence["value"] = sentence["value"].replace(DEFAULT_SOUND_TOKEN, f"{DEFAULT_SOUND_TOKEN}\n")
+                sentence["value"] = sentence["value"].replace(f"{DEFAULT_SOUND_TOKEN}\n\n", f"{DEFAULT_SOUND_TOKEN}\n")
+            if DEFAULT_SPEECH_TOKEN in sentence["value"]:
+                sentence["value"] = sentence["value"].replace(DEFAULT_SPEECH_TOKEN, f"{DEFAULT_SPEECH_TOKEN}\n")
+                sentence["value"] = sentence["value"].replace(f"{DEFAULT_SPEECH_TOKEN}\n\n", f"{DEFAULT_SPEECH_TOKEN}\n")
+    return sources
+
+
+def preprocess_plain(
+    sources: Sequence[str],
+    tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    # add end signal and concatenate together
+    conversations = []
+    for source in sources:
+        assert len(source) == 2
+        assert DEFAULT_IMAGE_TOKEN in source[0]["value"]
+        source[0]["value"] = DEFAULT_IMAGE_TOKEN
+        conversation = source[0]["value"] + source[1]["value"] + conversation_lib.default_conversation.sep
+        conversations.append(conversation)
+    # tokenize conversations
+    input_ids = [tokenizer_image_token(prompt, tokenizer, return_tensors="pt") for prompt in conversations]
+    targets = copy.deepcopy(input_ids)
+    for target, source in zip(targets, sources):
+        tokenized_len = len(tokenizer_image_token(source[0]["value"], tokenizer))
+        target[:tokenized_len] = IGNORE_INDEX
+
+    return dict(input_ids=input_ids, labels=targets)
+
+
+def preprocess(
+    sources: Sequence[str],
+    tokenizer: transformers.PreTrainedTokenizer,
+    has_image: bool = False,
+    no_system_prompt: bool = False,
+) -> Dict:
+    if conversation_lib.default_conversation.sep_style == conversation_lib.SeparatorStyle.PLAIN:
+        return preprocess_plain(sources, tokenizer)
+    return default_collate(
+        [
+            preprocess_conversation(conversation, tokenizer, no_system_prompt=no_system_prompt)
+            for conversation in sources
+        ]
+    )
+
+
+class LazySupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning.
+    This class is originally implemented by the LLaVA team and modified by
+    Ji Lin and Haotian Tang.
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        image_folder: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        training_args: TrainingArguments,
+    ):
+        super().__init__()
+        try:
+            with open(data_path) as fp:
+                list_data_dict = json.load(fp)
+        except:
+            with open(data_path) as fp:
+                list_data_dict = [json.loads(q) for q in fp]
+
+        # rank0_print("Formatting inputs...Skip in lazy mode")
+        print("Formatting inputs...Skip in lazy mode")
+        self.tokenizer = tokenizer
+        self.list_data_dict = list_data_dict
+        self.data_args = data_args
+        self.image_folder = image_folder
+        self.wav_processor = AutoFeatureExtractor.from_pretrained('/lustre/fsw/portfolios/adlr/users/sreyang/flamingo_v2/NV-Whisper')
+
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    @property
+    def lengths(self):
+        length_list = []
+        for sample in self.list_data_dict:
+            img_tokens = 128 if "image" in sample else 0
+            length_list.append(sum(len(conv["value"].split()) for conv in sample["conversations"]) + img_tokens)
+        return length_list
+
+    @property
+    def modality_lengths(self):
+        length_list = []
+        for sample in self.list_data_dict:
+            if 'duration' in sample.keys():
+                duration = sample["duration"]
+            else:
+                duration = 10.
+            try:
+                cur_len = sum(len(conv["value"].split()) for conv in sample["conversations"]) + int(math.ceil(duration * 25))
+                cur_len = cur_len if "sound" in sample else -cur_len
+                length_list.append(cur_len)
+            except:
+                try:
+                    cur_len = 0 + int(math.ceil(duration * 25))
+                    cur_len = cur_len if "sound" in sample else -cur_len
+                    length_list.append(cur_len)  
+                except:
+                    cur_len = 0 + int(math.ceil(10. * 25))
+                    cur_len = cur_len if "sound" in sample else -cur_len
+                    length_list.append(cur_len) 
+        return length_list
+    
+    @staticmethod
+    def _load_sound(sound_file, wav_processor, sample_rate=16000, window_length=30.0, window_overlap=0.0, max_num_window=3, audio_start = 0.0):
+        if sound_file is None:
+            return None
+        window_length  = int(window_length * sample_rate)
+        window_overlap = int(window_overlap * sample_rate)
+        max_num_window = int(max_num_window)
+        duration = max_num_window * (window_length - window_overlap) + window_overlap
+
+        sound_outputs = []
+        audio_feature_masks = []
+        audio_embed_masks = []
+
+        try:
+            sound_filename = str.split(sound_file, '/')[-1]
+            if '.ark' in sound_filename:
+                sound = kaldiio.load_mat(sound_file)
+                audio_data = sound[1]
+                audio_data=audio_data.astype(np.float16)
+            else:
+                audio_data = load_audio(sound_file, sample_rate, duration, audio_start) # already cuts to max duration
+            T = len(audio_data)
+            audio_data = audio_data.reshape(1, -1)
+            num_windows, full_length = get_num_windows(T, sample_rate, max_num_window)
+
+            audio_data_tensor = torch.from_numpy(int16_to_float32(float32_to_int16(audio_data))).float()
+            
+            for i in range(num_windows):
+                audio_embed_mask = torch.zeros(750)
+                start = i * (window_length - window_overlap)
+                audio_data_tensor_this = audio_data_tensor[:, start:start+window_length]
+                orig_length = audio_data_tensor_this.shape[1]
+                audio_data_tensor_this = wav_processor(audio_data_tensor_this.cpu().numpy(), sampling_rate=sample_rate, return_tensors="pt") #.squeeze(0) text="dummy", audios=audio_data_tensor_this, return_tensors="pt") #
+                sound_outputs.append(audio_data_tensor_this["input_features"])
+                # calculate the mask for the input melspec to Whisper
+                melspec_frames_this_window = int(math.ceil(orig_length / 160))
+                feature_attention_mask = torch.zeros(3000, dtype=torch.int32)
+                feature_attention_mask[:melspec_frames_this_window] = 1
+                audio_feature_masks.append(feature_attention_mask.unsqueeze(0))
+                # calculate the mask for the output embedding for use in AF2
+                conv_lengths = (melspec_frames_this_window - 1) // 2 + 1
+                output_embedding_lengths = (conv_lengths - 2) // 2 + 1
+                audio_embed_mask[:output_embedding_lengths] = 1
+                audio_embed_masks.append(audio_embed_mask)
+        except:
+            print('error loading file', sound_file)
+            sound_outputs.append(torch.zeros(1,128,3000))
+            audio_feature_masks.append(torch.zeros(1,3000, dtype=torch.int32))
+            audio_embed_masks.append(torch.zeros(750))
+
+        return torch.stack(sound_outputs, dim=0), torch.stack(audio_feature_masks, dim=0), torch.stack(audio_embed_masks, dim=0)
+
+    @staticmethod
+    def _load_speech(speech_path,sample_rate=16000):
+        if speech_path is None:
+            return None
+
+        speech_outputs = []
+        try:
+            speech = whisper.load_audio(speech_path)
+            speech = whisper.pad_or_trim(speech)
+            mel = whisper.log_mel_spectrogram(speech)
+            speech_outputs.append(mel.unsqueeze(0))
+        except:
+            speech_outputs.append(torch.zeros(1,80,3000))
+        return torch.stack(speech_outputs, dim=0)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        sources = self.list_data_dict[i]
+        if isinstance(i, int):
+            sources = [sources]
+        assert len(sources) == 1, "Don't know why it is wrapped to a list"  # FIXME
+        
+        import re
+        if "sound" in self.list_data_dict[i]:
+            # chat data loading 
+            if isinstance(self.list_data_dict[i]["sound"],list):
+                sound_files = self.list_data_dict[i]["sound"]
+                conversations_raw = self.list_data_dict[i]["conversations"]
+
+                # Step 1: Extract <sound-X> tags in order of appearance
+                sound_tag_pattern = re.compile(r"<sound-(\d+)>")
+                ordered_sound_tags = []
+
+                for turn in conversations_raw:
+                    tags = sound_tag_pattern.findall(turn["value"])
+                    ordered_sound_tags.extend([f"<sound-{tag}>" for tag in tags])
+
+                # Step 2: Load sound tensors in the order of tags
+                sound_tensor = []
+                audio_feature_masks = []
+                audio_embed_masks = []
+                sound_token_map = {}
+
+                for tag in ordered_sound_tags:
+                    idx = int(tag.split('-')[1][:-1])
+                    if tag not in sound_token_map:
+                        this_sound_tensor, af_mask, ae_mask = self._load_sound(sound_file, self.wav_processor, max_num_window=self.data_args.audio_frames)
+                        this_sound_tensor = this_sound_tensor.squeeze(1)  # (windows x 750 x 2048)
+                        sound_token_map[tag] = ("<sound>\n" * this_sound_tensor.shape[0]).rstrip()
+                        sound_tensor.append(this_sound_tensor)
+                        audio_feature_masks.append(af_mask)
+                        audio_embed_masks.append(ae_mask)
+                    else:
+                        # If already loaded, still append to match sequence
+                        this_sound_tensor, af_mask, ae_mask = self._load_sound(sound_file, self.wav_processor, max_num_window=self.data_args.audio_frames)
+                        this_sound_tensor = this_sound_tensor.squeeze(1)
+                        sound_tensor.append(this_sound_tensor)
+                        audio_feature_masks.append(af_mask)
+                        audio_embed_masks.append(ae_mask)
+                   
+
+                # Process conversations and inject sound markers
+                conversation = []
+                for turn in conversations_raw:
+                    role = turn["from"]
+                    value = turn["value"]
+
+                    # Replace any <sound-X> tag with corresponding repeated <sound>\n
+                    for tag, sound_token in sound_token_map.items():
+                        value = value.replace(tag, sound_token)
+
+                    conversation.append({
+                        "from": role,
+                        "value": value.rstrip()
+                    })
+
+                sources = [conversation]
+                sound_tensor = torch.cat(sound_tensor, dim=0)
+                audio_feature_masks = torch.cat(audio_feature_masks, dim=0)
+                audio_embed_masks = torch.cat(audio_embed_masks, dim=0)
+            else:
+                sound_file = self.list_data_dict[i]["sound"]
+                question = str(self.list_data_dict[i]["conversations"][0]["value"].rstrip())
+                answer = str(self.list_data_dict[i]["conversations"][1]["value"]).rstrip()
+                question = question.replace("<speech>\n", "").replace("\n<speech>", "").replace("<speech>", "")
+                question = question.replace("<sound>\n", "").replace("\n<sound>", "").replace("<sound>", "")
+                question = question.replace("<en><asr>\n", "").replace("\n<en><asr>", "").replace("<en><asr>", "")
+                question = question.replace("<eng><asr>\n", "").replace("\n<eng><asr>", "").replace("<eng><asr>", "")
+                sound_tensor, audio_feature_masks, audio_embed_masks = self._load_sound(sound_file, self.wav_processor, max_num_window=self.data_args.audio_frames)
+                sound_tensor=sound_tensor.squeeze(1) # squeeze the irrelevant dimension which was caused due to processor getting 1 batch for processing --> (windows x 750 x 2048)
+                question = "<sound>\n" * sound_tensor.shape[0] + question
+                conversation = [
+                    {"from": "human", "value": question},
+                    {"from": "gpt", "value": answer},
+                ]
+
+                sources = [conversation]    
+        data_dict = preprocess(
+            sources,
+            self.tokenizer,
+            has_image=(
+                "sound" in self.list_data_dict[i]
+            ),
+        )
+        if isinstance(i, int):
+            data_dict = dict(input_ids=data_dict["input_ids"][0], labels=data_dict["labels"][0])
+        
+        if "sound" in self.list_data_dict[i]:
+            data_dict["sound"] = sound_tensor
+            data_dict["sound_feature_masks"] = audio_feature_masks
+            data_dict["sound_embed_masks"] = audio_embed_masks
+        if "speech" in self.list_data_dict[i]:
+            data_dict["speech"] = speech_tensor
+      
+        return data_dict
+
+
+class LazyMMC4Dataset(Dataset):
+    """Dataset for supervised fine-tuning.
+    This class is implemented by Ji Lin and Haotian Tang."""
+
+    def __init__(
+        self,
+        data_path: str,
+        image_folder: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        training_args: TrainingArguments,
+        image_following_text_only=False,
+        text_only=False,
+    ):
+        super().__init__()
+
+        import pickle
+
+        n_samples = []
+        # actually shards and stats info
+        n_shards = len(os.listdir(data_path)) // 2
+        # n_shards = 100
+        count_info_list = sorted([f for f in os.listdir(data_path) if f.endswith(".count")])[:n_shards]
+        n_samples = [int(open(os.path.join(data_path, f)).read().strip()) for f in count_info_list]
+
+        print("total MMC4 samples", sum(n_samples))  # 10,881,869
+
+        PROCESS_GROUP_MANAGER = get_pg_manager()
+        if PROCESS_GROUP_MANAGER is not None:
+            import torch.distributed as dist
+
+            sequence_parallel_size = training_args.seq_parallel_size
+        else:
+            sequence_parallel_size = 1
+        print("sequence_parallel_size", sequence_parallel_size)
+        rank = training_args.process_index // sequence_parallel_size  # int(os.environ["RANK"])
+        world_size = training_args.world_size // sequence_parallel_size  # int(os.environ["WORLD_SIZE"])
+        shared_size = n_shards // world_size
+
+        gpu_samples = [sum(n_samples[i * shared_size : (i + 1) * shared_size]) for i in range(world_size)]
+        self.n_samples = min(gpu_samples) * world_size  # total size
+        self.idx_offset = rank * min(gpu_samples)
+        shard_start, shard_end = rank * shared_size, (rank + 1) * shared_size
+        print(f" * loading data from shard {shard_start}-{shard_end}")
+
+        shard_names = [d.replace(".count", ".pkl") for d in count_info_list]
+        shard_names = shard_names[shard_start:shard_end]
+
+        full_data_list = []
+        # now load data
+        for shard_name in shard_names:
+            # load shard
+            with open(os.path.join(data_path, shard_name), "rb") as f:
+                data_list = pickle.load(f)
+
+            full_data_list.extend(data_list)
+
+        print(f"* loaded totally {len(full_data_list)} samples")
+
+        self.data_list = full_data_list
+
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+        self.image_folder = image_folder
+
+        self.image_following_text_only = image_following_text_only
+        self.text_only = text_only
+
+    def __len__(self):
+        # return len(self.data_list)
+        return self.n_samples
+
+    @property
+    def modality_lengths(self):
+        # Estimate the number of tokens after tokenization, used for length-grouped sampling
+        length_list = []
+        for info in self.data_list:
+            num_images = min(6, len(info["image_info"]))
+            sentences = [info["text_list"][x["matched_text_index"]] for x in info["image_info"][:num_images]]
+            # The unit of cur_len is "words". We assume 1 word = 2 tokens.
+            cur_len = num_images * self.num_image_tokens // 2 + sum([len(x) for x in sentences])
+            length_list.append(cur_len)
+        return length_list
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        info = self.data_list[i - self.idx_offset]
+
+        sentences = info["text_list"]
+        # kentang-mit@: remove existing <image> tokens in the sentences
+        for ix in range(len(sentences)):
+            # if this is an html tag, we still preserve its semantic meaning
+            sentences[ix] = sentences[ix].replace("<image>", "<IMAGE>")
+        sim_matrix = info["similarity_matrix"]  # we do not use this...
+
+        # convert images from base64 to PIL and filter based on image-text similarity
+        images, sentence_ixs = [], []
+        if not self.text_only:
+            for sample_image, sim_vec in zip(info["image_info"], sim_matrix):
+                image_base64 = sample_image["image_base64"]
+                rawbytes = base64.b64decode(image_base64)
+
+                sim_ix = sample_image["matched_text_index"]
+                # sim_ix = np.argmax(sim_vec)
+                # sim_score = sim_vec[sim_ix]
+
+                # filter to images >= 5KB
+                # if len(rawbytes) // 1000 <= 5:
+                #     continue
+                # if sim_score < 0.24:
+                #     continue
+                image = Image.open(io.BytesIO(rawbytes)).convert("RGB")
+
+                images.append(image)
+                sentence_ixs.append(sim_ix)
+
+        # constrain max num 6 images
+        max_num_images = 6
+        if len(images) > max_num_images:
+            images = images[:max_num_images]
+            sentence_ixs = sentence_ixs[:max_num_images]
+
+        # reorder images according to text insertion
+        images = [images[iii] for iii in np.argsort(sentence_ixs)]
+
+        # preprocess and tokenize text
+        for ix in sentence_ixs:
+            sentences[ix] = f"<image>\n{sentences[ix]}"
+
+        if self.image_following_text_only:
+            # use pad tokens to divide sentence pieces
+            text = self.tokenizer.pad_token.join(sentences)
+        else:
+            text = " ".join(sentences)
+        # whitespace cleanup
+        text = text.replace("<image> ", "<image>").replace(" <image>", "<image>")
+        text = f"{text}{self.tokenizer.eos_token}"  # add eos token
+
+        if len(images) > 0:
+            if self.data_args.image_aspect_ratio == "dynamic_s2":
+                images, block_sizes = dynamic_s2_process_images_and_prompt(
+                    images, text, self.data_args, self.image_folder
+                )
+            elif self.data_args.image_aspect_ratio == "dynamic":
+                images, text = dynamic_process_images_and_prompt(
+                    images, text, self.data_args, self.image_folder, max_tiles=6
+                )
+            else:
+                images = torch.stack([process_image(image, self.data_args, self.image_folder) for image in images])
+
+            # the same size for all images, so we concat
+            # cur_token_len = (
+            #     images[0].shape[-2] // self.multimodal_cfg["patch_size"]
+            # ) * (images[0].shape[-1] // self.multimodal_cfg["patch_size"])
+            # cur_token_len += self.multimodal_cfg["n_extra_patch"]
+        else:
+            images = None
+            # cur_token_len = 0
+
+        input_ids = tokenizer_image_token(
+            text,
+            self.tokenizer,
+            return_tensors="pt",
+        )
+
+        image_token_id = self.tokenizer.media_token_ids["image"]
+
+        # now check the case where the last token is image patch token
+        if input_ids[-1] == image_token_id:  # need to remove one last image
+            last_non_im_patch_indices = torch.where(input_ids != image_token_id)[0][-1] + 1
+            input_ids = input_ids[:last_non_im_patch_indices]
+
+        n_im_patch = (input_ids == image_token_id).sum().item()
+
+        if self.data_args.image_aspect_ratio != "dynamic_s2":
+            images = images[:n_im_patch]
+            assert len(images) == n_im_patch, print(text, input_ids)
+        assert len(input_ids.shape) == 1, "Unexpected shape of 'input_ids' from MMC4."
+        input_ids = (
+            torch.concat([torch.tensor([self.tokenizer.bos_token_id]), input_ids])
+            if self.tokenizer.bos_token_id is not None and input_ids[0] != self.tokenizer.bos_token_id
+            else input_ids
+        )
+        targets = input_ids.clone()
+
+        if self.image_following_text_only:  # keep only text after leading image token
+            # remove loss for any token before the first <image> token
+            label_idx = 0
+            while label_idx < targets.shape[-1] and targets[label_idx] != image_token_id:
+                targets[label_idx] = IGNORE_INDEX
+                label_idx += 1
+
+            pad_token = self.tokenizer.convert_tokens_to_ids([self.tokenizer.pad_token])[0]
+
+            pad_token_idxs = torch.where(targets == pad_token)[0]
+            for pad_token_idx in pad_token_idxs:
+                token_idx = pad_token_idx + 1
+                while token_idx < targets.shape[-1] and targets[token_idx] != image_token_id:
+                    targets[token_idx] = IGNORE_INDEX
+                    token_idx += 1
+            # do not train on padding tokens
+            targets[targets == pad_token] = IGNORE_INDEX
+
+        # mask image tokens is unnecessary for llava-1.5
+        # targets[targets == IMAGE_TOKEN_INDEX] = IGNORE_INDEX
+        # print(input_ids.shape)
+
+        data_dict = dict(input_ids=input_ids, labels=targets, image=images)
+        if self.data_args.image_aspect_ratio == "dynamic_s2":
+            data_dict["block_sizes"] = block_sizes
+
+        return data_dict
+
+
+class LazyCoyoDataset(Dataset):
+    """Dataset for supervised fine-tuning.
+    This class is implemented by Ji Lin and Haotian Tang."""
+
+    num_image_tokens = 576
+
+    def __init__(
+        self,
+        data_path: str,
+        image_folder: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        training_args: TrainingArguments,
+        # kentang-mit@: balance the total number of tokens for Coyo and MMC4.
+        n_samples_per_idx=4,
+    ):
+        super().__init__()
+
+        import pickle
+
+        n_samples = []
+        # actually shards and stats info
+        n_shards = len(os.listdir(data_path)) // 2
+        # n_shards = 100
+        count_info_list = sorted([f for f in os.listdir(data_path) if f.endswith(".count")])[:n_shards]
+        n_samples = [int(open(os.path.join(data_path, f)).read().strip()) for f in count_info_list]
+
+        print("total COYO samples", sum(n_samples))
+
+        PROCESS_GROUP_MANAGER = get_pg_manager()
+        if PROCESS_GROUP_MANAGER is not None:
+            import torch.distributed as dist
+
+            sequence_parallel_size = training_args.seq_parallel_size
+        else:
+            sequence_parallel_size = 1
+        print("sequence_parallel_size", sequence_parallel_size)
+        rank = training_args.process_index // sequence_parallel_size  # int(os.environ["RANK"])
+        world_size = training_args.world_size // sequence_parallel_size  # int(os.environ["WORLD_SIZE"])
+        shared_size = n_shards // world_size
+
+        gpu_samples = [
+            sum(n_samples[i * shared_size : (i + 1) * shared_size]) // n_samples_per_idx for i in range(world_size)
+        ]
+        self.n_samples = min(gpu_samples) * world_size  # total size
+        self.idx_offset = rank * min(gpu_samples)
+
+        shard_start, shard_end = rank * shared_size, (rank + 1) * shared_size
+        print(f" * loading data from shard {shard_start}-{shard_end}")
+
+        shard_names = [d.replace(".count", ".pkl") for d in count_info_list]
+        shard_names = shard_names[shard_start:shard_end]
+
+        full_data_list = []
+        # now load data
+        for shard_name in shard_names:
+            # load shard
+            with open(os.path.join(data_path, shard_name), "rb") as f:
+                shard_data = pickle.load(f)
+                random.seed(42)
+                if "mmc4" in data_path:
+                    random.shuffle(shard_data)  # shuffle for MMC4cap only
+                full_data_list.extend(shard_data)
+
+        print(f"* loaded totally {len(full_data_list)} samples")
+
+        # now pack the samples into groups
+        n_groups = len(full_data_list) // n_samples_per_idx
+        full_data_list = [
+            full_data_list[i : i + n_samples_per_idx] for i in range(0, len(full_data_list), n_samples_per_idx)
+        ]
+        if len(full_data_list[-1]) < n_samples_per_idx:
+            full_data_list = full_data_list[:-1]
+        assert len(full_data_list) == n_groups
+        print(f"split into {n_groups} groups")
+
+        self.data_list = full_data_list
+
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+        self.image_folder = image_folder
+
+    def __len__(self):
+        # return len(self.data_list)
+        return self.n_samples
+
+    @property
+    def modality_lengths(self):
+        # Estimate the number of tokens after tokenization, used for length-grouped sampling
+        length_list = []
+        for samples in self.data_list:
+            cur_len = sum([len(conv["text" if "text" in conv else "caption"].split()) for conv in samples])
+            # The unit of cur_len is "words". We assume 1 word = 2 tokens.
+            cur_len = cur_len + len(samples) * self.num_image_tokens // 2
+            length_list.append(cur_len)
+        return length_list
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        CONCAT_SAMPLES = False
+        info_list = self.data_list[i - self.idx_offset]
+
+        text_list = []
+        image_list = []
+
+        for sample in info_list:
+            caption_key = (
+                "text" if "text" in sample else "caption"
+            )  # kentang-mit@: remove existing <image> tokens in the sentences
+            # kentang-mit@: remove existing <image> token.
+            # if this is an html tag, we still preserve its semantic meaning
+            sample[caption_key] = sample[caption_key].replace("<image>", "<IMAGE>")
+            text_list.append(DEFAULT_IMAGE_TOKEN + "\n" + sample[caption_key] + self.tokenizer.eos_token)
+            if "image" in sample:
+                image_base64 = sample["image"]
+                rawbytes = base64.b64decode(image_base64)
+            else:
+                rawbytes = sample["rawbytes"]
+            image = Image.open(io.BytesIO(rawbytes)).convert("RGB")
+            image_list.append(image)
+
+        image_list = torch.stack([process_image(image, self.data_args, self.image_folder) for image in image_list])
+
+        if CONCAT_SAMPLES:
+            # into <image>cap<eos><image>cap<eos>...
+            text_list = "".join(text_list)
+
+            input_ids = self.tokenizer(
+                text_list,
+                return_tensors="pt",
+                padding="longest",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+            ).input_ids  # 4, seq_len
+
+            input_ids = input_ids[0]
+
+        else:
+            input_ids = [
+                tokenizer_image_token(
+                    prompt,
+                    self.tokenizer,
+                    return_tensors="pt",
+                )
+                for prompt in text_list
+            ]
+            # print([x.shape[0] for x in input_ids], [len(x.split()) for x in text_list], [len(re.findall(r"<image[^>]*>", x)) for x in text_list])
+
+            # input_ids = torch.nn.utils.rnn.pad_sequence(
+            #     input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+            # )
+
+        targets = copy.deepcopy(input_ids)
+        for i in range(len(targets)):
+            targets[i][targets[i] == self.tokenizer.pad_token_id] = IGNORE_INDEX
+
+        return dict(input_ids=input_ids, labels=targets, image=image_list)
+
+
+class LazyWDSDataset(Dataset):
+    """Dataset for supervised fine-tuning.
+    This class is implemented by Ji Lin and Ligeng Zhu."""
+
+    def __init__(
+        self,
+        data_path: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        image_folder: str,
+        training_args: TrainingArguments,
+    ):
+        super().__init__()
+        n_samples = []
+        n_shards = len(os.listdir(data_path)) // 3
+        for shard in range(n_shards):
+            with open(os.path.join(data_path, f"{shard:05d}_stats.json")) as f:
+                info = json.load(f)
+                n_samples.append(info["successes"])
+
+        # print(f"[DEBUG] {data_path} total samples", sum(n_samples))  # 10,881,869
+
+        PROCESS_GROUP_MANAGER = get_pg_manager()
+        if PROCESS_GROUP_MANAGER is not None:
+            import torch.distributed as dist
+
+            sequence_parallel_size = training_args.seq_parallel_size
+        else:
+            sequence_parallel_size = 1
+        print("sequence_parallel_size", sequence_parallel_size)
+        rank = training_args.process_index // sequence_parallel_size  # int(os.environ["RANK"])
+        world_size = training_args.world_size // sequence_parallel_size  # int(os.environ["WORLD_SIZE"])
+        shared_size = n_shards // world_size
+        print("rank", rank, "world_size", world_size, "shared_size", shared_size)
+        gpu_samples = [sum(n_samples[i * shared_size : (i + 1) * shared_size]) for i in range(world_size)]
+        self.n_samples = min(gpu_samples) * world_size  # total size
+        self.idx_offset = rank * min(gpu_samples)
+        shard_start, shard_end = rank * shared_size, (rank + 1) * shared_size
+        print(f" * loading data from shard {shard_start}-{shard_end}")
+
+        tar_list = [f"{shard_idx:05d}.tar" for shard_idx in range(shard_start, shard_end)]
+
+        self.data_list = []
+        t1 = time.time()
+        for tar in tar_list:
+            tmp_path = f"/tmp/ccs{tar}"
+            tar_path = os.path.join(data_path, tar)
+
+            if PROCESS_GROUP_MANAGER is not None:
+                dist.barrier()
+                if PROCESS_GROUP_MANAGER.sp_rank == 0:
+                    os.makedirs(tmp_path, exist_ok=True)
+                    os.system(f"tar -xkf {tar_path} -C {tmp_path}")
+                dist.barrier()
+            else:
+                os.makedirs(tmp_path, exist_ok=True)
+                os.system(f"tar -xkf {tar_path} -C {tmp_path}")
+
+            txt_list = [f for f in os.listdir(tmp_path) if f.endswith(".txt")]
+
+            for txt in txt_list:
+                caption = open(os.path.join(tmp_path, txt)).read().strip()
+                image_path = os.path.join(tmp_path, txt.split(".")[0] + ".jpg")
+                self.data_list.append({"caption": caption, "image": image_path})
+        t2 = time.time()
+        print(f"Loading done. Total time: {t2 - t1:.2f} seconds")
+
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+        self.image_folder = image_folder
+
+    def __len__(self):
+        return self.n_samples
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+
+        # print("i", i, "idx_offset", self.idx_offset, "len", len(self.data_list))
+        info = self.data_list[i - self.idx_offset]
+        caption, image_path = info["caption"], info["image"]
+
+        rand_prompt = "<image>\n"
+        sources = [
+            {
+                "image": image_path,
+                "conversations": [
+                    {"from": "human", "value": rand_prompt},
+                    {"from": "gpt", "value": caption},
+                ],
+            }
+        ]
+
+        # one example of sources
+        # [{'id': 'GCC_train_001738742', 'image': 'GCC_train_001738742.jpg', 'conversations': [{'from': 'human', 'value': 'Provide a brief description of the given image.\n<image>'}, {'from': 'gpt', 'value': 'a sketch of an ostrich'}]}]
+        if "image" in sources[0]:
+            image = process_image(sources[0]["image"], self.data_args, self.image_folder)
+            image = torch.unsqueeze(image, dim=0)
+            # now random pick some context samples for training
+            if hasattr(self.data_args, "num_shots"):
+                if self.data_args.num_shots > 0:
+                    raise NotImplementedError
+        else:
+            raise NotImplementedError
+
+        data_dict = preprocess([sources[0]["conversations"]], self.tokenizer, has_image=True)
+
+        if isinstance(i, int):
+            data_dict = dict(input_ids=data_dict["input_ids"][0], labels=data_dict["labels"][0])
+
+        # image exist in the data
+        if image is not None:
+            data_dict["image"] = image
+        else:
+            raise NotImplementedError
+
+        return data_dict
+
+
+class LazyCCSWebDataset(Dataset):
+    """Dataset for supervised fine-tuning.
+    This class is implemented by Ligeng Zhu."""
+
+    def __init__(
+        self,
+        data_path: str,
+        image_folder: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        training_args: TrainingArguments,
+    ):
+        super().__init__()
+        t1 = time.time()
+
+        from llava.data.simple_vila_webdataset import VILAWebDataset
+
+        print("[DEBUG] ", osp.abspath(data_path))
+        self.dataset = VILAWebDataset(data_path=osp.abspath(data_path))
+
+        t2 = time.time()
+        print(f"Loading done. Total time: {t2 - t1:.2f} seconds")
+
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        # info = self.data_list[i - self.idx_offset]
+        # caption, image_path = info["caption"], info["image"]
+        info = self.dataset[i]
+        if ".jpg" in info:
+            caption, image_path = info[".txt"], info[".jpg"]
+        elif ".png" in info:
+            caption, image_path = info[".txt"], info[".png"]
+        elif ".webp" in info:
+            caption, image_path = info[".txt"], info[".webp"]
+        elif ".bmp" in info:
+            caption, image_path = info[".txt"], info[".bmp"]
+        elif ".tiff" in info:
+            caption, image_path = info[".txt"], info[".tiff"]
+        else:
+            print(info.keys())
+            print(info)
+            raise KeyError
+
+        caption = caption.replace("<image>", "<IMAGE>")
+        if isinstance(image_path, io.BytesIO):
+            image_path = Image.open(image_path).convert("RGB")
+
+        if not isinstance(image_path, PIL.Image.Image):
+            print(image_path)
+            print(info.keys())
+            print(type(image_path))
+            raise NotImplementedError
+
+        rand_prompt = "<image>\n"
+        sources = [
+            {
+                "image": image_path,
+                "conversations": [
+                    {"from": "human", "value": rand_prompt},
+                    {"from": "gpt", "value": caption},
+                ],
+            }
+        ]
+
+        # one example of sources
+        # [{'id': 'GCC_train_001738742', 'image': 'GCC_train_001738742.jpg', 'conversations': [{'from': 'human', 'value': 'Provide a brief description of the given image.\n<image>'}, {'from': 'gpt', 'value': 'a sketch of an ostrich'}]}]
+        if "image" in sources[0]:
+            image = process_image(sources[0]["image"], self.data_args, image_folder=None)
+            image = torch.unsqueeze(image, dim=0)
+            # now random pick some context samples for training
+            if hasattr(self.data_args, "num_shots"):
+                if self.data_args.num_shots > 0:
+                    raise NotImplementedError
+        else:
+            raise NotImplementedError
+
+        data_dict = preprocess([sources[0]["conversations"]], self.tokenizer, has_image=True)
+
+        if isinstance(i, int):
+            data_dict = dict(input_ids=data_dict["input_ids"][0], labels=data_dict["labels"][0])
+
+        # image exist in the data
+        if image is not None:
+            data_dict["image"] = image
+        else:
+            raise NotImplementedError
+
+        return data_dict
+
+
+from functools import lru_cache
+
+
+@lru_cache(maxsize=16)
+def lru_json_load(fpath):
+    with open(fpath) as fp:
+        return json.load(fp)
+
+
+class LazyCoyoWebDataset(Dataset):
+    """Dataset for supervised fine-tuning.
+    This class is implemented by Ligeng Zhu."""
+
+    num_image_tokens = 576
+
+    def __init__(
+        self,
+        data_path: str,
+        image_folder: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        training_args: TrainingArguments,
+        # kentang-mit@: balance the total number of tokens for Coyo and MMC4.
+        n_samples_per_idx=4,
+    ):
+        super().__init__()
+
+        from llava.data.simple_vila_webdataset import VILAWebDataset
+
+        print("[DEBUG] ", osp.abspath(data_path))
+        self.dataset = VILAWebDataset(data_path=osp.abspath(data_path), meta_path=data_args.meta_path)
+
+        if data_args.start_idx >= 0 and data_args.end_idx >= 0:
+            # Ligeng: support slicing for ablate different subsets.
+            total = len(self.dataset)
+            start_idx = int(total * data_args.start_idx)
+            end_idx = int(total * data_args.end_idx)
+            print(f"loading subset from {start_idx} to {end_idx}, total {total}")
+            self.dataset = torch.utils.data.Subset(self.dataset, range(start_idx, end_idx))
+
+        # For caption choice,
+        #   if None: use original caption
+        #   if a folder path: use specified caption to override original one (choice1)
+        #   if a folder path: use specified caption and concat with original one (choice2)
+        self.caption_choice = None
+        self.caption_choice_2 = None
+        self.data_path = data_path
+
+        if data_args.caption_choice is not None:
+            self.caption_choice = data_args.caption_choice
+            print("[recap] Override coyo caption using ", self.caption_choice)
+
+        if data_args.caption_choice_2 is not None:
+            self.caption_choice_2 = data_args.caption_choice_2
+            print("[recapv2] Override coyo caption using ", self.caption_choice_2)
+
+        print("total samples", len(self.dataset))
+        PROCESS_GROUP_MANAGER = get_pg_manager()
+        if PROCESS_GROUP_MANAGER is not None:
+            import torch.distributed as dist
+
+            sequence_parallel_size = training_args.seq_parallel_size
+            sequence_parallel_rank = PROCESS_GROUP_MANAGER.sp_rank
+        else:
+            sequence_parallel_size = 1
+        print("sequence_parallel_size", sequence_parallel_size)
+        rank = (
+            training_args.process_index // sequence_parallel_size if "RANK" in os.environ else 2
+        )  # int(os.environ["RANK"])
+        world_size = (
+            training_args.world_size // sequence_parallel_size if "WORLD_SIZE" in os.environ else 32
+        )  # int(os.environ["WORLD_SIZE"])
+        print(
+            "rank",
+            rank,
+            "world_size",
+            world_size,
+        )
+
+        self.n_samples_per_idx = n_samples_per_idx
+        # self.n_samples = len(self.dataset) // n_samples_per_idx
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+
+    def __len__(self):
+        return len(self.dataset) // self.n_samples_per_idx
+
+    @property
+    def modality_lengths(self):
+        # Estimate the number of tokens after tokenization, used for length-grouped sampling
+        length_list = []
+        for samples in self.data_list:
+            cur_len = sum([len(conv["text" if "text" in conv else "caption"].split()) for conv in samples])
+            # The unit of cur_len is "words". We assume 1 word = 2 tokens.
+            cur_len = cur_len + len(samples) * self.num_image_tokens // 2
+            length_list.append(cur_len)
+        return length_list
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        CONCAT_SAMPLES = False
+        # info_list = self.dataset[i - self.idx_offset]
+
+        begin_idx, end_idx = (
+            i * self.n_samples_per_idx,
+            (i + 1) * self.n_samples_per_idx,
+        )
+        end_idx = min(end_idx, len(self.dataset))
+
+        text_list = []
+        image_list = []
+
+        for idx in range(begin_idx, end_idx):
+            info = self.dataset[idx]
+            if ".jpg" in info:
+                caption, image_path = info[".txt"], info[".jpg"]
+            elif ".png" in info:
+                caption, image_path = info[".txt"], info[".png"]
+            elif ".webp" in info:
+                caption, image_path = info[".txt"], info[".webp"]
+            elif ".bmp" in info:
+                caption, image_path = info[".txt"], info[".bmp"]
+            elif ".tiff" in info:
+                caption, image_path = info[".txt"], info[".tiff"]
+            else:
+                print(info.keys())
+                print(info)
+                raise KeyError
+
+            if self.caption_choice is not None:
+                # load new captions
+                shard = info["__shard__"]
+                url = info[".json"]["url"]
+                tar_name = osp.relpath(osp.realpath(shard), osp.realpath(self.data_path))
+                # tar_name = osp.dirname(shard)
+                shard_json_path = osp.join(self.caption_choice, tar_name + ".json")
+                try:
+                    shard_json = lru_json_load(shard_json_path)
+                    try:
+                        caption = shard_json[url]["output"]
+                    except KeyError:
+                        print(f"{url} not in caption. fallback to original caption temporarially")
+                except:
+                    print(f"shard_json_path {shard_json_path} not found. fallback to original caption temporarially")
+            caption = caption.replace("<image>", "<IMAGE>")
+            text_list.append(DEFAULT_IMAGE_TOKEN + caption + self.tokenizer.eos_token)
+
+            if isinstance(image_path, io.BytesIO):
+                image_path = Image.open(image_path).convert("RGB")
+
+            if not isinstance(image_path, PIL.Image.Image):
+                print(image_path)
+                print(info.keys())
+                print(type(image_path))
+                raise NotImplementedError
+
+            image_list.append(image_path)
+
+        # image_list = torch.stack([process_image(image, self.data_args, image_folder=None) for image in image_list])
+        # NOTE(fix by ligeng)
+        #  now image_list should return a list of image tensor where each has a dimension of (1, c, h, w)
+        image_list = [process_image(image, self.data_args, image_folder=None).unsqueeze(0) for image in image_list]
+
+        if CONCAT_SAMPLES:
+            # into <image>cap<eos><image>cap<eos>...
+            text_list = "".join(text_list)
+
+            input_ids = self.tokenizer(
+                text_list,
+                return_tensors="pt",
+                padding="longest",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+            ).input_ids  # 4, seq_len
+
+            input_ids = input_ids[0]
+        else:
+            input_ids = [
+                tokenizer_image_token(
+                    prompt,
+                    self.tokenizer,
+                    return_tensors="pt",
+                )
+                for prompt in text_list
+            ]
+            input_ids = [
+                (
+                    torch.concat([torch.tensor([self.tokenizer.bos_token_id]), input_ids_i])
+                    if input_ids_i[0] != self.tokenizer.bos_token_id
+                    else input_ids_i
+                )
+                for input_ids_i in input_ids
+            ]
+
+        targets = copy.deepcopy(input_ids)
+        for i in range(len(targets)):
+            targets[i][targets[i] == self.tokenizer.pad_token_id] = IGNORE_INDEX
+
+        return dict(input_ids=input_ids, labels=targets, image=image_list)
+
+
+class LazyVideoWebDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(
+        self,
+        data_path: str,
+        image_folder: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        training_args: TrainingArguments,
+        # cache_path: str,
+        # n_samples_per_idx=4,
+    ):
+        super().__init__()
+
+        # from llava.data.simple_video_dataset import SimpleVideoDataset
+
+        from llava.data.simple_vila_webdataset import VILAWebDataset
+
+        print("[DEBUG] ", osp.abspath(data_path))
+        self.dataset = VILAWebDataset(
+            data_path=osp.abspath(data_path),
+            meta_path=f"{osp.abspath(data_path)}/wids-meta.json",
+            # cache_dir=cache_path,
+        )
+
+        # None: use original caption
+        # Folder path: use original caption
+        self.caption_choice = None
+        self.data_path = data_path
+
+        if data_args.caption_choice is not None:
+            self.caption_choice = data_args.caption_choice
+            print("[recap] Override LazyVideo caption using ", self.caption_choice)
+
+        print("total samples", len(self.dataset))
+        # InternVid: TODO
+        PROCESS_GROUP_MANAGER = get_pg_manager()
+        if PROCESS_GROUP_MANAGER is not None:
+            import torch.distributed as dist
+
+            sequence_parallel_size = training_args.seq_parallel_size
+            sequence_parallel_rank = PROCESS_GROUP_MANAGER.sp_rank
+        else:
+            sequence_parallel_size = 1
+        print("sequence_parallel_size", sequence_parallel_size)
+        rank = (
+            training_args.process_index // sequence_parallel_size if "RANK" in os.environ else 2
+        )  # int(os.environ["RANK"])
+        world_size = (
+            training_args.world_size // sequence_parallel_size if "WORLD_SIZE" in os.environ else 32
+        )  # int(os.environ["WORLD_SIZE"])
+        print(
+            "rank",
+            rank,
+            "world_size",
+            world_size,
+        )
+        self.rank = rank
+        # rank = int(os.environ["RANK"]) if "RANK" in os.environ else 2
+        # world_size = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 32
+
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+
+        self.missing_uids = set()
+
+    def __len__(self):
+        return len(self.dataset)
+
+    @property
+    def modality_lengths(self):
+        # Estimate the number of tokens after tokenization, used for length-grouped sampling
+        length_list = []
+        for samples in self.data_list:
+            cur_len = sum([len(conv["text" if "text" in conv else "caption"].split()) for conv in samples])
+            # The unit of cur_len is "words". We assume 1 word = 2 tokens.
+            cur_len = cur_len + len(samples) * self.num_image_tokens // 2
+            length_list.append(cur_len)
+        return length_list
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        ADD_TEXT_PROMPT = False
+        num_video_frames = self.data_args.num_video_frames if hasattr(self.data_args, "num_video_frames") else 8
+        loader_fps = self.data_args.fps if hasattr(self.data_args, "fps") else 0.0
+
+        info = self.dataset[i]
+
+        caption = ""
+        # print(info)
+        if ".mp4" in info:
+            caption, video_path = info[".txt"], info[".mp4"]
+        else:
+            video_path = None
+            caption = "Empty video."
+
+        images, frames_loaded, _ = LazySupervisedDataset._load_video(
+            video_path, num_video_frames, loader_fps, self.data_args
+        )
+
+        if frames_loaded == 0:
+            caption = "Empty video."
+
+        if self.caption_choice is not None:
+            shard = info["__shard__"]
+            uuid = osp.join(info["__shard__"], info["__key__"])
+            url = info["__key__"]
+            tar_name = osp.basename(info["__shard__"])
+
+            try:
+                shard_json_path = osp.join(self.caption_choice, tar_name.replace(".tar", ".json"))
+                shard_json = lru_json_load(shard_json_path)
+                caption = shard_json[url]["summary"]["output"]
+            except (KeyError, FileNotFoundError, json.decoder.JSONDecodeError):
+                if uuid not in self.missing_uids:
+                    print("override caption not found for ", uuid)
+                    self.missing_uids.add(uuid)
+
+            # print(f"[DEBUG {uuid}]", caption)
+
+        frames_loaded_successfully = len(images)
+        if caption is None:
+            caption = ""
+        prompt = "<image>\n" * frames_loaded_successfully + caption
+        image_tensor = torch.stack([process_image(image, self.data_args, None) for image in images])
+
+        input_ids = tokenizer_image_token(
+            prompt,
+            self.tokenizer,
+            return_tensors="pt",
+        )
+        targets = copy.deepcopy(input_ids)
+        data_dict = dict(input_ids=input_ids, labels=targets, image=image_tensor)
+
+        return data_dict
+
+
+class DataCollatorForSupervisedDatasetSeqParallel:
+    """Collate examples for supervised fine-tuning.
+    This class is originally implemented by the LLaVA team and
+    modified by Haotian Tang."""
+
+    def __init__(
+        self,
+        tokenizer: transformers.PreTrainedTokenizer,
+        data_args: DataArguments,
+        training_args: TrainingArguments,
+        sp_degree: int,
+        sp_rank: int,
+        ring_degree: int,
+        ring_type: str,
+    ):
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+        self.training_args = training_args
+        self.sp_degree = sp_degree
+        self.sp_rank = sp_rank
+        self.ring_degree = ring_degree
+        self.ring_type = ring_type
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels, images = [], [], []
+        image_token_id = self.tokenizer.media_token_ids["image"]
+        video_token_id = self.tokenizer.media_token_ids["video"]
+
+        for instance in instances:
+            if not isinstance(instance["input_ids"], list):
+                input_ids.append(instance["input_ids"])
+            else:
+                input_ids += instance["input_ids"]
+            if not isinstance(instance["labels"], list):
+                labels.append(instance["labels"])
+            else:
+                labels += instance["labels"]
+            # Note (kentang-mit@: we do not directly push tensors to
+            # images, but list of tensors.
+            if "video" in instance:
+                instance["image"] = torch.cat(instance["video"])
+                video_id_pos = torch.where(input_ids[-1] == video_token_id)[0][0]
+                replace_ids = torch.Tensor(
+                    ([image_token_id] + self.tokenizer.encode("\n")) * instance["image"].shape[0],
+                    device=input_ids[-1].device,
+                )
+                input_ids[-1] = torch.cat(
+                    [input_ids[-1][:video_id_pos], replace_ids, input_ids[-1][video_id_pos + 1 :]]
+                ).to(input_ids[-1].dtype)
+                labels[-1] = torch.cat(
+                    [
+                        labels[-1][:video_id_pos],
+                        torch.Tensor([IGNORE_INDEX] * instance["image"].shape[0] * 2),
+                        labels[-1][video_id_pos + 1 :],
+                    ]
+                ).to(labels[-1].dtype)
+                instance.pop("video")
+
+            if "image" in instance:
+                cur_image = instance["image"]
+                assert len(cur_image.shape) == 4
+                # n_images, 3, size, size
+                if cur_image.shape[0] == 0:
+                    warnings.warn("loaded one sample without images.")
+                if not isinstance(instance["input_ids"], list):
+                    # datasets other than coyo, not packing >1 samples together
+                    images.append(cur_image)
+                else:
+                    # coyo-like datasets
+                    images.extend(cur_image.chunk(cur_image.size(0), 0))
+            else:
+                warnings.warn("loaded one sample without images.")
+                images.append([])
+        # kentang-mit@: we need to make sure these two lists have
+        # the same length. We will use input_ids to filter out images corresponding
+        # to truncated <image> tokens later.
+
+        max_num_images = max([len(_images) for _images in images])
+        for _images, _input_ids in zip(images, input_ids):
+            assert (
+                len(_images) == (_input_ids == image_token_id).sum().item()
+            ), f"Number mismatch between images and placeholder image tokens in 'len(_images) == (_input_ids == image_token_id).sum().item()'.\
+                Expect to have {len(_images)} images but only found {(_input_ids == image_token_id).sum().item()} images in tokens. \
+                Error input_ids: {_input_ids} {self.tokenizer.decode([x if x != -200 else 200 for x in _input_ids])}"
+
+        NUM_TOKENS_PER_IMAGE = self.data_args.num_image_tokens
+        if hasattr(self.data_args.image_processor, "crop_size"):
+            crop_size = self.data_args.image_processor.crop_size
+        else:
+            crop_size = self.data_args.image_processor.size
+
+        # Init the padding sample
+        seq_id = 0
+        while seq_id < len(input_ids):
+            # Skip the samples without images
+            dummy_image = torch.ones((1, 3, crop_size["height"], crop_size["width"]), device=input_ids[seq_id].device)
+            # dummy input_ids include one bos, one image token, and one eos
+            dummy_input_ids = torch.zeros_like(input_ids[seq_id][:3])
+            dummy_input_ids[0] = self.tokenizer.bos_token_id
+            dummy_input_ids[1] = image_token_id
+            dummy_input_ids[2] = self.tokenizer.eos_token_id
+            dummy_labels = copy.deepcopy(dummy_input_ids)
+            dummy_labels[:2] = IGNORE_INDEX
+            dummy_seqlen = NUM_TOKENS_PER_IMAGE + 2  # TODO: Check the hard coding of 2
+            dummy_position_ids = torch.arange(start=0, end=dummy_seqlen, dtype=torch.int32)
+            break
+
+        # Sort with the real length of the sequence
+        combined = sorted(
+            zip(input_ids, labels, images),
+            key=lambda x: len(x[2]) * (NUM_TOKENS_PER_IMAGE - 1) + x[0].size(-1),
+            reverse=True,  # Start Packing from the sequence with most images.
+        )
+        sorted_ids, sorted_labels, sorted_images = zip(*combined)
+        sorted_ids, sorted_labels, sorted_images = list(sorted_ids), list(sorted_labels), list(sorted_images)
+        max_seq_length = self.tokenizer.model_max_length  # len(sorted_ids[0])
+        max_sample_len = 0
+
+        batches = []
+        label_batches = []
+        position_ids = []
+        batch_images = []
+        seqlens_in_batch = []
+
+        i = 0
+        while i < len(sorted_ids):
+            current_batch = torch.tensor([], dtype=torch.int32)
+            current_label_batch = torch.tensor([], dtype=torch.int32)
+            current_position_ids = torch.tensor([], dtype=torch.int32)
+            current_batch_images = []
+            current_num_images = 0
+            current_len = 0
+            current_num_samples = 0
+
+            # Pack a few samples into one sample
+            while i < len(sorted_ids):
+                num_images = (sorted_ids[i] == image_token_id).sum().item()
+                num_image_tokens_added = num_images * (NUM_TOKENS_PER_IMAGE - 1)
+                num_incoming_tokens = sorted_ids[i].size(-1) + num_image_tokens_added
+
+                # Handle RingAttn_Varlen which requires `seqlens_in_batch` should be divisible by `ring_degree`
+                if self.ring_degree > 1:
+                    RING_PAD_TOKEN_INDEX = 2
+                    if self.ring_type == "ring_varlen":
+                        if num_incoming_tokens % self.sp_degree != 0:
+                            pad_len = self.sp_degree - num_incoming_tokens % self.sp_degree
+                            num_incoming_tokens += pad_len
+                            # pad `input_ids`
+                            pad_tensor = torch.full(
+                                (pad_len,), RING_PAD_TOKEN_INDEX, dtype=sorted_ids[i].dtype, device=sorted_ids[i].device
+                            )
+                            sorted_ids[i] = torch.cat([sorted_ids[i], pad_tensor])
+
+                            # pad `label`
+                            pad_label_tensor = torch.full(
+                                (pad_len,), IGNORE_INDEX, dtype=sorted_labels[i].dtype, device=sorted_labels[i].device
+                            )
+                            sorted_labels[i] = torch.cat([sorted_labels[i], pad_label_tensor])
+                    elif self.ring_type == "zigzag_ring_varlen":
+                        self.zigzag_sp_degree = self.sp_degree * 2
+                        if num_incoming_tokens % self.zigzag_sp_degree != 0:
+                            pad_len = self.zigzag_sp_degree - num_incoming_tokens % self.zigzag_sp_degree
+                            num_incoming_tokens += pad_len
+                            # pad `input_ids`
+                            pad_tensor = torch.full(
+                                (pad_len,), RING_PAD_TOKEN_INDEX, dtype=sorted_ids[i].dtype, device=sorted_ids[i].device
+                            )
+                            sorted_ids[i] = torch.cat([sorted_ids[i], pad_tensor])
+
+                            # pad `label`
+                            pad_label_tensor = torch.full(
+                                (pad_len,), IGNORE_INDEX, dtype=sorted_labels[i].dtype, device=sorted_labels[i].device
+                            )
+                            sorted_labels[i] = torch.cat([sorted_labels[i], pad_label_tensor])
+                    else:
+                        raise ValueError(f"Invalid ring_type: {self.ring_type}")
+
+                if num_incoming_tokens > max_seq_length:
+                    print(
+                        f"Warning: Skipping one packed sample with {num_incoming_tokens} tokens,\
+                        please consider increase max seq len {max_seq_length}."
+                    )
+                    i += 1
+                    continue
+
+                if (
+                    (current_num_images == 0)
+                    or (current_num_images < self.sp_degree)
+                    or (
+                        (current_num_images + num_images <= max_num_images)
+                        and (current_len + num_incoming_tokens <= max_sample_len)
+                    )
+                ) and (current_len + num_incoming_tokens <= max_seq_length):
+                    current_num_images += num_images
+                    current_len += num_incoming_tokens
+                    current_num_samples += 1
+                    current_position_ids = torch.cat(
+                        (current_position_ids, torch.arange(start=0, end=num_incoming_tokens)), dim=0
+                    )
+                    current_batch = torch.cat((current_batch, sorted_ids[i]), dim=0)
+                    sorted_labels[i][0] = IGNORE_INDEX
+                    current_label_batch = torch.cat((current_label_batch, sorted_labels[i]), dim=0)
+                    seqlens_in_batch.append(num_incoming_tokens)
+                    current_batch_images.extend(sorted_images[i])
+                    i += 1
+                    assert current_num_images == len(current_batch_images)
+                else:
+                    break
+
+            # Padding the sample with the dummy image sample, if there are no enough images
+            MAX_RETRY = self.sp_degree
+            num_retry = 0
+            while current_num_images < self.sp_degree and current_len < max_seq_length and num_retry <= MAX_RETRY:
+                current_num_images += dummy_image.size(0)
+                current_len += dummy_seqlen
+                current_num_samples += 1
+                current_position_ids = torch.cat((current_position_ids, dummy_position_ids), dim=0)
+                current_batch = torch.cat((current_batch, dummy_input_ids), dim=0)
+                current_label_batch = torch.cat((current_label_batch, dummy_labels), dim=0)
+                seqlens_in_batch.append(dummy_seqlen)
+                current_batch_images.extend(dummy_image)
+                # We pad from left side to ensure correct grad flow
+                # current_batch = torch.cat((dummy_input_ids, current_batch), dim=0)
+                # current_label_batch = torch.cat((dummy_labels, current_label_batch), dim=0)
+                # seqlens_in_batch.insert(0, dummy_seqlen)
+                # current_batch_images = torch.cat((dummy_image, current_batch_images), dim=0)
+                num_retry += 1
+
+            # Drop the samples that do not have enough images
+            if current_num_images < self.sp_degree:
+                print(f"Warning: Skipping one packed sample with {current_num_images} images")
+                seqlens_in_batch = seqlens_in_batch[:-current_num_samples]
+                continue
+
+            max_sample_len = max(max_sample_len, current_len)
+            batches.append(current_batch)
+            label_batches.append(current_label_batch)
+            position_ids.append(current_position_ids)
+            batch_images.append(current_batch_images)
+
+            try:
+                assert current_num_images == len(torch.where(current_batch == image_token_id)[0].tolist())
+            except AssertionError:
+                print(f"Error num_images on {self.sp_rank}", current_num_images)
+                print("current_batch", current_batch)
+                print(
+                    f"Error len(torch.where(batches[i] == image_token_id)[0].tolist() on {self.sp_rank}:",
+                    len(torch.where(current_batch == image_token_id)[0].tolist()),
+                )
+                print(f"Error len(current_batch_images) on {self.sp_rank}:", len(current_batch_images))
+                raise AssertionError
+
+        # Split for sequence parallelism
+        for i in range(len(batches)):
+            image_token_indices = torch.where(batches[i] == image_token_id)[0].tolist()
+            image_ids = torch.arange(0, len(image_token_indices), dtype=torch.int32)
+            batches[i] = extract_local_input_ids(
+                batches[i], image_token_indices, self.sp_rank, self.sp_degree, self.tokenizer.bos_token_id
+            )
+            label_batches[i] = extract_local_input_ids(
+                label_batches[i], image_token_indices, self.sp_rank, self.sp_degree, self.tokenizer.bos_token_id
+            )
+            batch_images[i] = torch.concat(
+                extract_local_from_list(batch_images[i], self.sp_rank, self.sp_degree), dim=0
+            )
+            H, W = batch_images[i].size(-2), batch_images[i].size(-1)
+            batch_images[i] = batch_images[i].reshape(-1, 3, W, H)
+            num_images = len(batch_images[i])
+
+            try:
+                assert num_images == len(torch.where(batches[i] == image_token_id)[0].tolist())
+            except AssertionError:
+                print(f"Error num_images on {self.sp_rank}", num_images)
+                print("batches[i]", batches[i])
+                print(
+                    f"Error len(torch.where(batches[i] == image_token_id)[0].tolist() on {self.sp_rank}:",
+                    len(torch.where(batches[i] == image_token_id)[0].tolist()),
+                )
+                print(f"Error batch_images[i] on {self.sp_rank}:", batch_images[i].shape)
+                raise AssertionError
+            position_ids[i] = extract_local_position_ids(
+                position_ids[i], image_token_indices, image_ids, self.sp_rank, self.sp_degree, NUM_TOKENS_PER_IMAGE - 1
+            )
+
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            batches, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(label_batches, batch_first=True, padding_value=IGNORE_INDEX)
+        seqlens_in_batch = [torch.tensor(x) for x in seqlens_in_batch]
+        seqlens_in_batch = torch.stack(seqlens_in_batch, axis=0)
+        seqlens_in_batch = seqlens_in_batch.flatten()
+        position_ids = torch.nn.utils.rnn.pad_sequence(position_ids, batch_first=True, padding_value=-1)
+
+        if batch_images:
+            batch_images = [torch.unbind(images) for images in batch_images]
+            flat_batch_images = [item for sublist in batch_images for item in sublist]
+        else:
+            flat_batch_images = None
+        batch = dict(
+            input_ids=input_ids,
+            labels=labels,
+            # notice that we inject attention mask here
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+            seqlens_in_batch=seqlens_in_batch,
+            media={"image": flat_batch_images},
+            media_config={"image": {}},
+            position_ids=position_ids,
+        )
+        return batch
+
+
+def make_supervised_data_module(
+    tokenizer: PreTrainedTokenizer,
+    data_args: DataArguments,
+    training_args: TrainingArguments,
+) -> Dict:
+    """Make dataset and collator for supervised fine-tuning.
+    This function is originally implemented by the LLaVA team and
+    modified by Jason Lu, Haotian Tang and Ligeng Zhu."""
+    datasets_mixture.register_datasets_mixtures()
+
+    from .builder import build_dataset
+
+    train_dataset = build_dataset(data_args.data_mixture, data_args, training_args, tokenizer)
+    training_args.sample_lens = [len(d) for d in train_dataset.datasets]
+
+    PROCESS_GROUP_MANAGER = get_pg_manager()
+    if PROCESS_GROUP_MANAGER is None:
+        data_collator = DataCollator(tokenizer=tokenizer)
+    else:
+        sp_degree = training_args.seq_parallel_size
+        sp_rank = PROCESS_GROUP_MANAGER.sp_rank
+        ring_degree = PROCESS_GROUP_MANAGER.ring_degree
+        ring_type = PROCESS_GROUP_MANAGER.ring_type
+        data_collator = DataCollatorForSupervisedDatasetSeqParallel(
+            tokenizer=tokenizer,
+            data_args=data_args,
+            training_args=training_args,
+            sp_degree=sp_degree,
+            sp_rank=sp_rank,
+            ring_degree=ring_degree,
+            ring_type=ring_type,
+        )
+
+    return dict(
+        train_dataset=train_dataset,
+        data_collator=data_collator,
+    )

llava/data/datasets_mixture.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import warnings
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Dataset:
+    dataset_name: str
+    dataset_type: str = field(default="torch")
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+    meta_path: str = field(default=None, metadata={"help": "Path to the meta data for webdataset."})
+    image_path: str = field(default=None, metadata={"help": "Path to the training image data."})
+    speech_path: str = field(default=None, metadata={"help": "Path to the training speech data."})
+    caption_choice: str = field(default=None, metadata={"help": "Path to the caption directory for recaption."})
+    description: str = field(
+        default=None,
+        metadata={
+            "help": "Detailed desciption of where the data is from, how it is labelled, intended use case and the size of the dataset."
+        },
+    )
+    test_script: str = (None,)
+    maintainer: str = (None,)
+    ############## ############## ############## ############## ############## ##############
+    caption_choice: str = field(default=None, metadata={"help": "Path to the captions for webdataset."})
+    caption_choice_2: str = field(default=None, metadata={"help": "Path to the captions for webdataset."})
+    start_idx: float = field(default=-1, metadata={"help": "Start index of the dataset."})
+    end_idx: float = field(default=-1, metadata={"help": "Start index of the dataset."})
+
+
+DATASETS_LEGACY = {}
+
+
+def add_dataset(dataset):
+    if dataset.dataset_name in DATASETS_LEGACY:
+        # make sure the data_name is unique
+        warnings.warn(f"{dataset.dataset_name} already existed in DATASETS. Make sure the name is unique.")
+    assert "+" not in dataset.dataset_name, "Dataset name cannot include symbol '+'."
+    DATASETS_LEGACY.update({dataset.dataset_name: dataset})
+
+
+def register_datasets_mixtures():
+    ############## ############## ############## ############## ############## ############## 
+    # Audio Datasets 
+    ############## ############## ############## ############## ############## ##############
+    
+    data_mixture_1 = Dataset(
+        dataset_name="data_mixture_1",
+        dataset_type="torch",
+        data_path="/path/to/your/data_mixture_1/train.json",
+    )
+    add_dataset(data_mixture_1)
+
+    data_mixture_2 = Dataset(
+        dataset_name="data_mixture_2",
+        dataset_type="torch",
+        data_path="/path/to/your/data_mixture_2/train.json",
+    )
+    add_dataset(data_mixture_2)
+    # Add more data mixtures below

llava/data/registry/datasets/audio_test.yaml

@@ -0,0 +1,97 @@
+---
+Clotho-AQA-AQA:
+  _target_: llava.data.LLaVADataset
+  data_path: Clotho-AQA-AQA/test.json
+Music-AVQA-AQA_All:
+  _target_: llava.data.LLaVADataset
+  data_path: Music-AVQA-AQA_All/test.json
+CochlScene-SceneClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: CochlScene-SceneClassification/test.json
+NSynth-Source:
+  _target_: llava.data.LLaVADataset
+  data_path: NSynth-Source/test.json
+NSynth-Instrument:
+  _target_: llava.data.LLaVADataset
+  data_path: NSynth-Instrument/test.json
+FSD50k-EventClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: FSD50k-EventClassification/test.json
+Clotho-v2-AudioCaptioning:
+  _target_: llava.data.LLaVADataset
+  data_path: Clotho-v2-AudioCaptioning/test.json
+audiocaps-AudioCaptioning:
+  _target_: llava.data.LLaVADataset
+  data_path: audiocaps-AudioCaptioning/test.json
+ravdess-EmotionClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: ravdess-EmotionClassification/val.json
+GTZAN-GenreClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: GTZAN-GenreClassification/test.json
+UrbanSound8K-EventClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: UrbanSound8K-EventClassification/train.json
+Medley-solos-DB-InstrClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: Medley-solos-DB-InstrClassification/test.json
+ESC50-EventClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: ESC50-EventClassification/train.json
+CREMA-D-EmotionClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: CREMA-D-EmotionClassification/test.json
+IEMOCAP-EmotionClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: IEMOCAP-EmotionClassification/test.json
+MELD-EmotionClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: MELD-EmotionClassification/test.json
+MELD-SentimentClassification:
+  _target_: llava.data.LLaVADataset
+  data_path: MELD-SentimentClassification/test.json
+MMAU:
+  _target_: llava.data.LLaVADataset
+  data_path: MMAU/test.json
+MMAU-mini:
+  _target_: llava.data.LLaVADataset
+  data_path: MMAU/test-mini.json
+AudioEntailmentQA:
+  _target_: llava.data.LLaVADataset
+  data_path: AudioEntailmentQA/test.json
+SPGI-ASR:
+  _target_: llava.data.LLaVADataset
+  data_path: SPGI-ASR/val.json
+SWBD-ASR:
+  _target_: llava.data.LLaVADataset
+  data_path: SWBD-ASR/val.json
+LibriSpeech-ASR-clean:
+  _target_: llava.data.LLaVADataset
+  data_path: LibriSpeech-ASR/test_clean.json
+LibriSpeech-ASR-other:
+  _target_: llava.data.LLaVADataset
+  data_path: LibriSpeech-ASR/test_other.json
+VoxPopuli-ASR:
+  _target_: llava.data.LLaVADataset
+  data_path: VoxPopuli-ASR/test.json
+Europarl-ASR:
+  _target_: llava.data.LLaVADataset
+  data_path: Europarl-ASR/test.json
+CV-ASR:
+  _target_: llava.data.LLaVADataset
+  data_path: CV-ASR/test.json
+GigaSpeech-ASR:
+  _target_: llava.data.LLaVADataset
+  data_path: GigaSpeech-ASR/test.json
+CompA-R-AQA:
+  _target_: llava.data.LLaVADataset
+  data_path: CompA-R-AQA/test.json
+MuschoMusicQA:
+  _target_: llava.data.LLaVADataset
+  data_path: MuschoMusicQA/test.json
+CMM:
+  _target_: llava.data.LLaVADataset
+  data_path: CMM/test.json
+AIR-Bench:
+  _target_: llava.data.LLaVADataset
+  data_path: AIR-Bench/test.json

llava/data/registry/datasets/default.yaml

@@ -0,0 +1,5 @@
+---
+dummy:
+    _target_: llava.data.DummyDataset
+    num_instances: 10000
+    comments: dummy dataset for testing

llava/data/registry/mixtures.yaml

@@ -0,0 +1,78 @@
+---
+audio_speech_all:
+    -CV-ASR_1
+    -MELD-EmotionClassification+
+    -BBCSoundEffects-AudioDescription
+    -SWBD-ASR_1
+    -WavCaps-SoundBible-AudioCaptioning
+    -AudioSet-Speech-Audio-QA
+    -SONYC-UST-EventClassification
+    -VoxPopuli-ASR_1
+    -FSD50k-EventClassification
+    -SalmonnQA
+    -emov-db-EmotionClassification
+    -LLARK_MagnaTagATune-mir+tess-EmotionClassification
+    -Europarl-ASR_1
+    -jl-corpus-EmotionClassification
+    -Ego-10-AudioCaptioning
+    -SPGI-ASR_1
+    -CREMA-D-EmotionClassification
+    -MusicBenchQA
+    -WavCaps-BBC_Sound_Effects-AudioCaptioning
+    -NSynth-Instrument
+    -SpokenSquadQA
+    -NSynth-MIR
+    -AudioEntailmentQA
+    -GigaSpeech-ASR_1
+    -WavCaps-AudioSet_SL-AudioCaptioning
+    -NonSpeech7k-EventClassification
+    -chime-home-EventClassification
+    -MusicCaps-AudioCaptioning
+    -LP-MusicCaps-MSD-AudioCaptioning
+    -Ego-30-AudioCaptioning
+    -NSynth-Source+Clotho-v2-AudioCaptioning
+    -LP-MusicCaps-MC-AudioCaptioning
+    -Clotho-AQA-EventClassification
+    -WavCaps-FreeSound-AudioCaptioning
+    -LLARK_MagnaTagATune-reasoning
+    -AudioSet-Temporal-Speech-Audio-QA
+    -TUT-EventClassification
+    -ESC50-EventClassification
+    -WavText5K-Tagging
+    -MELD-SentimentClassification
+    -Music-AVQA-AQA_All
+    -Music-AVQA-AVQA_All
+    -MACS-AudioCaptioning
+    -Medley-solos-DB-InstrClassification
+    -AudioSet-EventClassification
+    -OMGEmotion-EmotionClassification
+    -FMA-GenreClassification
+    -Epidemic_sound-AudioCaptioning
+    -CochlScene-SceneClassification
+    -LLARK_FMA-reasoning
+    -ravdess-EmotionClassification
+    -CompA-R-AQA
+    -MU-LLAMA-AQA
+    -musdbhq-InstrClassification
+    -UrbanSound8K-EventClassification
+    -audiocaps-AudioCaptioning
+    -VocalSound-VocalClassification
+    -CLAP_freesound-AudioCaptioning
+    -MMAUQA
+    -SongDescriber-AudioCaptioning
+    -HeySQuADQA
+    -Mira-AudioCaptioning
+    -Clotho-AQA-AQA
+    -LibriSpeech-ASR_1
+    -IEMOCAP-EmotionClassification
+    -AudioSetFullwoAudioMusicCaps-EventClassification
+    -MSP-PODCAST-Publish-1.9-EmotionClassification
+    -OpenAQA-AQA
+    -SoundDescs-AudioDescription
+    -LibriSQA
+    -LLARK_FMA-mir
+    -LP-MusicCaps-MTT-AudioCaptioning
+    -GTZAN-GenreClassification
+    -musdbhq-captioning
+    -YesNoQA
+

llava/entry.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import typing
+from typing import List, Optional
+
+if typing.TYPE_CHECKING:
+    from transformers import PreTrainedModel
+else:
+    PreTrainedModel = None
+
+__all__ = ["load"]
+
+
+def load(
+    model_path: str,
+    model_base: Optional[str] = None,
+    devices: Optional[List[int]] = None,
+    **kwargs,
+) -> PreTrainedModel:
+    import torch
+
+    from llava.conversation import auto_set_conversation_mode
+    from llava.mm_utils import get_model_name_from_path
+    from llava.model.builder import load_pretrained_model
+
+    auto_set_conversation_mode(model_path)
+
+    model_name = get_model_name_from_path(model_path)
+    model_path = os.path.expanduser(model_path)
+    if os.path.exists(os.path.join(model_path, "model")):
+        model_path = os.path.join(model_path, "model")
+
+    # Set `max_memory` to constrain which GPUs to use
+    if devices is not None:
+        assert "max_memory" not in kwargs, "`max_memory` should not be set when `devices` is set"
+        kwargs.update(max_memory={device: torch.cuda.get_device_properties(device).total_memory for device in devices})
+
+    model = load_pretrained_model(model_path, model_name, model_base, **kwargs)[1]
+    return model

llava/eval/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import os
+
+from llava.utils import io
+
+__all__ = ["EVAL_ROOT", "TASKS"]
+
+
+EVAL_ROOT = "scripts/eval"
+TASKS = io.load(os.path.join(os.path.dirname(__file__), "registry_audio.yaml"))

llava/eval/eval_audio_bench.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import argparse
+import csv
+import itertools
+import json
+import os
+
+import torch
+from datasets import load_dataset
+from tqdm import tqdm
+
+import llava
+from llava import conversation as conversation_lib
+from llava.data.builder import DATASETS
+from llava.eval.mmmu_utils.eval_utils import parse_choice
+from llava.utils import distributed as dist
+from llava.utils import io
+from llava.utils.logging import logger
+
+
+def load_existing_ids(output_file):
+    if not os.path.exists(output_file):
+        return set(), []
+    try:
+        with open(output_file, "r") as f:
+            lines = f.readlines()
+            outputs = [json.loads(line) for line in lines]
+            processed_ids = {item["id"] for item in outputs}
+            return processed_ids, outputs
+    except Exception as e:
+        print(f"Error loading existing outputs: {e}")
+        return set(), []
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default=None)
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--task", type=str, default=None)
+    parser.add_argument("--conv-mode", type=str, default="auto")
+    parser.add_argument("--generation-config", type=json.loads)
+    parser.add_argument("--output-dir", type=str, default=None)
+    args = parser.parse_args()
+
+    # Set up distributed environment
+    dist.init()
+    devices = range(dist.local_rank(), torch.cuda.device_count(), dist.local_size())
+    torch.cuda.set_device(devices[0])
+
+    # Load stage 3 model with line 56
+    model = llava.load(args.model_base, model_base=None, devices=devices)
+    # Uncomment line 58-63 to load stage 3.5 model on top of stage 3 for thinking mode and long audio mode
+    # model = PeftModel.from_pretrained(
+    #         model,
+    #         args.model_path,
+    #         device_map="auto",
+    #         torch_dtype=torch.float16,
+    #     )
+    # Set up generation config
+    generation_config = model.default_generation_config
+    if args.generation_config is not None:
+        generation_config.update(**args.generation_config)
+
+    # Load data and chunk it
+    json_file = DATASETS[args.task]["data_path"]
+    instances = io.load(json_file)
+    instances = instances[dist.rank() :: dist.size()]
+
+    output_path = os.path.join(args.output_dir, f"outputs_{args.task}.jsonl")
+    processed_ids, outputs = load_existing_ids(output_path)
+
+    count = len(outputs)
+    # Run inference
+    new_outputs = []
+    for instance in tqdm(instances, disable=not dist.is_main()):
+        uuid = instance["id"]
+        sound_path = instance["sound"]
+
+        if sound_path in processed_ids:
+            continue  # Skip if already processed
+        sound = llava.Sound(sound_path)
+        conversations = instance["conversations"]
+        question = conversations[0]["value"]
+
+        response = model.generate_content([sound, question], generation_config=generation_config)
+       
+        print("response", response)
+
+        output = {"id": sound_path, "question": question, "gt_answer": conversations[1]["value"],  "pred": response}
+        new_outputs.append(output)
+        count = count +1
+        if count % 20 == 0:
+            # Gather and save outputs
+            if dist.size() > 1:
+                outputs_new = dist.gather(new_outputs, dst=0)
+                if dist.is_main():
+                    outputs_new = list(itertools.chain(*outputs_new))
+                    final_outputs = outputs + outputs_new
+                    io.save(os.path.join(args.output_dir, f"outputs_{args.task}.jsonl"), final_outputs)
+            else:
+                final_outputs = outputs + new_outputs
+                io.save(os.path.join(args.output_dir, f"outputs_{args.task}.jsonl"), final_outputs)
+    if dist.size() > 1:
+        new_outputs = dist.gather(new_outputs, dst=0)
+        if not dist.is_main():
+            return
+        new_outputs = list(itertools.chain(*new_outputs))
+        final_outputs = outputs + new_outputs
+    io.save(os.path.join(args.output_dir, "outputs_"+str(args.task)+".jsonl"), final_outputs)
+
+if __name__ == "__main__":
+    main()

llava/eval/mmmu_utils/eval_utils.py

@@ -0,0 +1,61 @@
+# This file is originated from the official MMMU codebase:
+# https://github.com/MMMU-Benchmark/MMMU
+
+import random
+
+import numpy as np
+
+
+def parse_choice(response, all_choices, index2ans=None):
+    """
+    Parse the prediction from the generated response.
+    Return the predicted index e.g., A, B, C, D.
+    """
+    for char in [",", ".", "!", "?", ";", ":", "'"]:
+        response = response.strip(char)
+    response = " " + response + " "  # add space to avoid partial match
+
+    index_ans = True
+    ans_with_brack = False
+    candidates = []
+    for choice in all_choices:  # e.g., (A) (B) (C) (D)
+        if f"({choice})" in response:
+            candidates.append(choice)
+            ans_with_brack = True
+
+    if len(candidates) == 0:
+        for choice in all_choices:  # e.g., A B C D
+            if f" {choice} " in response:
+                candidates.append(choice)
+
+    # if all above doesn't get candidates, check if the content is larger than 5 tokens and try to parse the example
+    if len(candidates) == 0 and len(response.split()) > 5 and index2ans is not None:
+        for index, ans in index2ans.items():
+            if ans.lower() in response.lower():
+                candidates.append(index)
+                index_ans = False  # it's content ans.
+
+    if len(candidates) == 0:  # still not get answer, randomly choose one.
+        pred_index = random.choice(all_choices)
+    elif len(candidates) > 1:
+        start_indexes = []
+        if index_ans:
+            if ans_with_brack:
+                for can in candidates:
+                    index = response.rfind(f"({can})")
+                    start_indexes.append(index)  # -1 will be ignored anyway
+                # start_indexes = [generated_response.index(f'({can})') for can in candidates]
+            else:
+                for can in candidates:
+                    index = response.rfind(f" {can} ")
+                    start_indexes.append(index)
+        else:
+            for can in candidates:
+                index = response.lower().rfind(index2ans[can].lower())
+                start_indexes.append(index)
+        # get the last one
+        pred_index = candidates[np.argmax(start_indexes)]
+    else:  # if only one candidate, use it.
+        pred_index = candidates[0]
+
+    return pred_index

llava/eval/registry_audio.yaml

@@ -0,0 +1,93 @@
+Clotho-AQA-AQA:
+  tags:
+  - local
+Music-AVQA-AQA_All:
+  tags:
+  - local
+CochlScene-SceneClassification:
+  tags:
+  - local
+NSynth-Source:
+  tags:
+  - local
+NSynth-Instrument:
+  tags:
+  - local
+FSD50k-EventClassification:
+  tags:
+  - local
+Clotho-v2-AudioCaptioning:
+  tags:
+  - local
+audiocaps-AudioCaptioning:
+  tags:
+  - local
+ravdess-EmotionClassification:
+  tags:
+  - local
+GTZAN-GenreClassification:
+  tags:
+  - local
+UrbanSound8K-EventClassification:
+  tags:
+  - local
+Medley-solos-DB-InstrClassification:
+  tags:
+  - local
+ESC50-EventClassification:
+  tags:
+  - local
+CREMA-D-EmotionClassification:
+  tags:
+  - local
+IEMOCAP-EmotionClassification:
+  tags:
+  - local
+MELD-EmotionClassification:
+  tags:
+  - local
+MELD-SentimentClassification:
+  tags:
+  - local
+MMAU:
+  tags:
+  - local
+AudioEntailmentQA:
+  tags:
+  - local
+SPGI-ASR:
+  tags:
+  - local
+SWBD-ASR:
+  tags:
+  - local
+LibriSpeech-ASR-clean:
+  tags:
+  - local
+LibriSpeech-ASR-other:
+  tags:
+  - local
+VoxPopuli-ASR:
+  tags:
+  - local
+Europarl-ASR:
+  tags:
+  - local
+CV-ASR:
+  tags:
+  - local
+GigaSpeech-ASR:
+  tags:
+  - local
+CompA-R-AQA:
+  tags:
+  - local
+MuschoMusicQA:
+  tags:
+  - local
+CMM:
+  tags:
+  - local
+AIR-Bench:
+  tags:
+  - local

llava/media.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["Media", "File", "Image", "Video", "Speech",  "Sound"]
+
+
+class Media:
+    pass
+
+
+class File(Media):
+    def __init__(self, path: str) -> None:
+        self.path = path
+
+
+class Image(File):
+    pass
+
+
+class Video(File):
+    pass
+
+
+class Speech(File):
+    pass
+
+class Sound(File):
+    pass

llava/mm_utils.py

@@ -0,0 +1,641 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# dynamic_preprocess and find_closest_aspect_ratio are referenced from https://github.com/OpenGVLab/InternVL
+
+import base64
+import os
+import tempfile
+from io import BytesIO
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import StoppingCriteria
+
+from pydub import AudioSegment
+from torchvision import transforms
+import soundfile as sf
+from librosa import resample as librosa_resample
+import whisper
+import random
+from pytorchvideo.data.clip_sampling import ConstantClipsPerVideoSampler, UniformClipSampler
+DEFAULT_AUDIO_FRAME_SHIFT_MS = 10  # in milliseconds
+def get_frame_from_vcap(vidcap, num_frames=10, max_fps=0.0, fps=None, frame_count=None, video_file_name=None):
+    import cv2
+
+    if fps == None or frame_count == None:
+        # if one of fps or frame_count is None, still recompute
+        fps = vidcap.get(cv2.CAP_PROP_FPS)
+        frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if fps == 0 or frame_count == 0:
+        print(f"Video file not found. return empty images. {video_file_name}")
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * num_frames, 0, [0.]
+
+    duration = frame_count / fps
+    frame_interval = frame_count // num_frames
+    if frame_interval == 0 and frame_count <= 1:
+        print(f"frame_interval is equal to 0. return empty image. {video_file_name}")
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * num_frames, 0, [0.]
+    # print("duration:", duration, "frames:", frame_count, "intervals:", frame_interval)
+
+    images = []
+    count = 0
+    success = True
+    frame_indices = np.linspace(0, frame_count - 1, num_frames, dtype=int)
+    frame_times = [frame / fps for frame in frame_indices]
+    while success:
+        # print("frame_count:", frame_count, "count:", count, "num_frames:", num_frames, "frame_interval:", frame_interval)
+        if frame_count >= num_frames:
+            success, frame = vidcap.read()
+            if count in frame_indices:
+                try:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                except BaseException:
+                    continue
+                if len(images) >= num_frames:
+                    return images, num_frames, frame_times
+            count += 1
+        else:
+            # Left padding frames if the video is not long enough
+            success, frame = vidcap.read()
+            if success:
+                try:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                except BaseException:
+                    continue
+                count += 1
+            else:
+                break
+    if len(images) == 0:
+        raise ValueError("Did not find enough frames in the video. return empty image.")
+
+    return images, len(images), frame_times
+
+
+def get_frame_from_vcap_with_fps(vidcap, num_frames=10, max_fps=0.0, fps=None, frame_count=None, video_file_name=None):
+    """
+    num_frames is the max number of frames the model can support.
+    frame_count is the number of frames in the input video.
+    max_fps is the max FPS of the model can support.
+    fps is the fps of the input video.
+    """
+
+    import random
+
+    import cv2
+
+    if fps == None or frame_count == None:
+        # if one of fps or frame_count is None, still recompute
+        fps = vidcap.get(cv2.CAP_PROP_FPS)
+        frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    if fps == 0 or frame_count == 0:
+        print(f"Video file not found. return empty images. {video_file_name}")
+        empty_video_frames = int(random.uniform(2, 8 * max_fps))
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * empty_video_frames, 0, [0.]
+
+    duration = frame_count / fps
+    # print("duration:", duration, "frames:", frame_count, "fps:", fps, "num_frames:", num_frames, "max_fps:", max_fps)
+    # If the video is too long (longer than max_fps and num_frames can support),
+    # we will use lower fps to sample frames.
+    if duration >= num_frames / max_fps:
+        frame_interval = frame_count // num_frames
+
+        # If the video is too short, we will skip the video if there is only one frame.
+        if frame_interval == 0 and frame_count <= 1:
+            print(f"frame_interval is equal to 0. return empty image. {video_file_name}")
+            empty_video_frames = int(random.uniform(2, 8 * max_fps))
+            return [
+                Image.new("RGB", (720, 720)),
+            ] * empty_video_frames, 0, [0.]
+
+        images = []
+        count = 0
+        success = True
+        frame_indices = np.linspace(0, frame_count - 1, num_frames, dtype=int)
+        frame_times = [frame / fps for frame in frame_indices]
+        while success:
+            if frame_count >= num_frames:
+                # success, frame = vidcap.read()
+                if count in frame_indices:
+                    success, frame = vidcap.read()
+                    try:
+                        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                        im_pil = Image.fromarray(img)
+                        images.append(im_pil)
+                    except:
+                        # print("Failed to read frame:", count)
+                        continue
+                    if len(images) >= num_frames:
+                        return images, num_frames, frame_times
+                else:
+                    success = vidcap.grab()
+                count += 1
+            else:
+                # Left padding frames if the video is not long enough
+                success, frame = vidcap.read()
+                if success:
+                    try:
+                        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                        im_pil = Image.fromarray(img)
+                        images.append(im_pil)
+                    except:
+                        # print("Failed to read frame:", count)
+                        continue
+                    count += 1
+                else:
+                    break
+    else:
+        frames_required = int(duration * max_fps)
+        frame_indices = np.linspace(0, frame_count - 1, frames_required, dtype=int)
+        if frames_required == 0:
+            print(f"frames_required is fewer than 2. Duration {duration}, return empty image.")
+            empty_video_frames = int(random.uniform(2, 8 * max_fps))
+            return [
+                Image.new("RGB", (720, 720)),
+            ] * empty_video_frames, 0, [0.]
+        elif frames_required == 1:
+            frame_indices = np.linspace(0, frame_count - 1, 2, dtype=int)
+        images = []
+        count = 0
+        looked = 0
+        success = True
+
+        while success:
+            success, frame = vidcap.read()
+            if success and (looked in frame_indices):
+                try:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                except:
+                    continue
+                count += 1
+            looked += 1
+    frame_times = [frame / fps for frame in frame_indices]
+
+    if len(images) == 0:
+        empty_video_frames = int(random.uniform(2, 8 * max_fps))
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * empty_video_frames, 0, [0.]
+    else:
+        return images, len(images), frame_times
+
+
+def opencv_extract_frames(vpath_or_bytesio, frames=6, max_fps=0.0, fps=None, frame_count=None):
+    """
+    Extract frames from a video using OpenCV.
+
+    Args:
+        vpath_or_bytesio (str or BytesIO): Path to the video file or BytesIO object containing the video.
+        frames (int): Number of frames to extract from the video.
+        fps (float): Frames per second of the video. If 0.0, the function will extract frames at equal intervals.
+
+    Returns:
+        list: List of PIL Images extracted from the video.
+
+    Raises:
+        NotImplementedError: If the type of `vpath_or_bytesio` is not supported.
+    """
+    import cv2
+    if isinstance(vpath_or_bytesio, str):
+        vidcap = cv2.VideoCapture(vpath_or_bytesio)
+        if max_fps > 0.0:
+            return get_frame_from_vcap_with_fps(
+                vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=vpath_or_bytesio
+            )
+        return get_frame_from_vcap(
+            vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=vpath_or_bytesio
+        )
+    elif isinstance(vpath_or_bytesio, (BytesIO,)):
+        # assuming mp4
+        with tempfile.NamedTemporaryFile(delete=True, suffix=".mp4") as temp_video:
+            temp_video.write(vpath_or_bytesio.read())
+            temp_video_name = temp_video.name
+            vidcap = cv2.VideoCapture(temp_video_name)
+            if max_fps > 0.0:
+                return get_frame_from_vcap_with_fps(
+                    vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=temp_video_name
+                )
+            return get_frame_from_vcap(
+                vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=temp_video_name
+            )
+    else:
+        raise NotImplementedError(type(vpath_or_bytesio))
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+    """
+    Expand the given PIL image to a square shape by adding padding.
+
+    Parameters:
+    - pil_img: The PIL image to be expanded.
+    - background_color: The color of the padding to be added.
+
+    Returns:
+    - The expanded PIL image.
+
+    If the image is already square, it is returned as is.
+    If the image is wider than it is tall, padding is added to the top and bottom.
+    If the image is taller than it is wide, padding is added to the left and right.
+    """
+    width, height = pil_img.size
+    if pil_img.mode == "L":
+        background_color = background_color[0]
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float("inf")
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+
+def dynamic_preprocess(image, min_num=1, max_num=12, image_size=384, use_thumbnail=True):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    # calculate the existing image aspect ratio
+    target_ratios = {
+        (i, j)
+        for n in range(min_num, max_num + 1)
+        for i in range(1, n + 1)
+        for j in range(1, n + 1)
+        if i * j <= max_num and i * j >= min_num
+    }
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size,
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
+
+def dynamic_s2_preprocess(image, s2_scales=[384, 768, 1152], max_num=12, image_size=384):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+    min_num = (s2_scales[-1] // s2_scales[0]) ** 2  # at least use number of tiles as the largest scale
+
+    processed_images = []
+
+    ##########################################################################################
+    ############# Add tiles for all but the last scale using fixed squre ratio ###############
+    ##########################################################################################
+
+    for scale in s2_scales[:-1]:
+        target_width = image_size * (scale // s2_scales[0])
+        target_height = image_size * (scale // s2_scales[0])
+        blocks = (scale // s2_scales[0]) ** 2
+
+        # resize the image
+        resized_img = image.resize((target_width, target_height))
+        for i in range(blocks):
+            box = (
+                (i % (target_width // image_size)) * image_size,
+                (i // (target_width // image_size)) * image_size,
+                ((i % (target_width // image_size)) + 1) * image_size,
+                ((i // (target_width // image_size)) + 1) * image_size,
+            )
+            # split the image
+            split_img = resized_img.crop(box)
+            processed_images.append(split_img)
+
+    ##########################################################################################
+    ################ Add tiles for the last scale using dynamic aspect ratio #################
+    ##########################################################################################
+
+    # calculate the existing image aspect ratio
+    target_ratios = {
+        (i, j)
+        for n in range(min_num, max_num + 1)
+        for i in range(1, n + 1)
+        for j in range(1, n + 1)
+        if i * j <= max_num and i * j >= min_num
+    }
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size,
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+
+    return processed_images, (target_aspect_ratio[1], target_aspect_ratio[0])
+
+
+
+def dynamic_s2_process_images_and_prompt(images, data_args, image_folder=None):
+    idx = 0
+    all_images = []
+    all_block_size = []
+    for img in images:
+        processed_images, block_size = process_image(img, data_args, image_folder, enable_dynamic_s2=True)
+        all_images.append(processed_images)
+        all_block_size.append(block_size)
+        idx += 2
+    if all_images:
+        all_images = torch.cat(all_images)
+    else:
+        all_images = None
+    return all_images, all_block_size
+
+
+def process_image(
+    image_file, data_args, image_folder, enable_dynamic_res=False, enable_dynamic_s2=False, max_tiles=None
+):
+    processor = data_args.image_processor
+    if isinstance(image_file, str):
+        if image_folder is not None:
+            image = Image.open(os.path.join(image_folder, image_file)).convert("RGB")
+        else:
+            image = Image.open(image_file).convert("RGB")
+    else:
+        # image is stored in bytearray
+        image = image_file
+    image = image.convert("RGB")
+    if hasattr(data_args.image_processor, "crop_size"):
+        # CLIP vision tower
+        crop_size = data_args.image_processor.crop_size
+    else:
+        # SIGLIP vision tower
+        assert hasattr(data_args.image_processor, "size")
+        crop_size = data_args.image_processor.size
+    if "dynamic_s2" in data_args.image_aspect_ratio and enable_dynamic_s2:
+        assert crop_size["height"] == crop_size["width"]
+        images, block_size = dynamic_s2_preprocess(
+            image, s2_scales=data_args.s2_scales, max_num=data_args.max_tiles, image_size=crop_size["height"]
+        )
+        images = [processor.preprocess(image, return_tensors="pt")["pixel_values"][0] for image in images]
+        return torch.stack(images), block_size
+    if "dynamic" in data_args.image_aspect_ratio and enable_dynamic_res:
+        assert crop_size["height"] == crop_size["width"]
+        if max_tiles is not None:
+            max_num = max_tiles
+        else:
+            max_num = data_args.max_tiles
+        images = dynamic_preprocess(image, min_num=data_args.min_tiles, max_num=max_num, image_size=crop_size["height"])
+        images = [processor.preprocess(image, return_tensors="pt")["pixel_values"][0] for image in images]
+        return torch.stack(images)
+
+    if data_args.image_aspect_ratio == "resize":
+        image = image.resize((crop_size["width"], crop_size["height"]))
+    if data_args.image_aspect_ratio == "pad":
+
+        def expand2square(pil_img, background_color):
+            width, height = pil_img.size
+            if width == height:
+                return pil_img
+            elif width > height:
+                result = Image.new(pil_img.mode, (width, width), background_color)
+                result.paste(pil_img, (0, (width - height) // 2))
+                return result
+            else:
+                result = Image.new(pil_img.mode, (height, height), background_color)
+                result.paste(pil_img, ((height - width) // 2, 0))
+                return result
+
+        image = expand2square(image, tuple(int(x * 255) for x in processor.image_mean))
+        image = processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
+    else:
+        # Using default behavior of the vision encoder
+        # For CLIP, default is central crop
+        # For Radio, default is central crop
+        # For Siglip, default is resize
+        # For InternVIT, default is resize
+        image = processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
+    return image
+
+def get_num_windows(T, sr, max_num_window=5):
+
+    window_length  = int(30.0 * sr)
+    window_overlap = int(0.0 * sr)
+    max_num_window = max_num_window
+
+    num_windows = 1
+    if T <= window_length:
+        num_windows = 1
+        full_length = window_length
+    elif T >= (max_num_window * window_length - (max_num_window - 1) * window_overlap):
+        num_windows = max_num_window
+        full_length = (max_num_window * window_length - (max_num_window - 1) * window_overlap)
+    else:
+        num_windows = 1 + int(np.ceil((T - window_length) / float(window_length - window_overlap)))
+        full_length = num_windows * window_length - (num_windows - 1) * window_overlap
+    
+    return num_windows, full_length
+
+def load_audio(file_path, target_sr=16000, duration=30.0, start=0.0):
+    if file_path.endswith('.mp3'):
+        audio = AudioSegment.from_file(file_path)
+        if len(audio) > (start + duration) * 1000:
+            audio = audio[start * 1000:(start + duration) * 1000]
+
+        if audio.frame_rate != target_sr:
+            audio = audio.set_frame_rate(target_sr)
+
+        if audio.channels > 1:
+            audio = audio.set_channels(1)
+        
+        data = np.array(audio.get_array_of_samples())
+        if audio.sample_width == 2:
+            data = data.astype(np.float32) / np.iinfo(np.int16).max
+        elif audio.sample_width == 4:
+            data = data.astype(np.float32) / np.iinfo(np.int32).max
+        else:
+            raise ValueError("Unsupported bit depth: {}".format(audio.sample_width))
+
+    else:
+        with sf.SoundFile(file_path) as audio:
+            original_sr = audio.samplerate
+            channels = audio.channels
+
+            max_frames = int((start + duration) * original_sr)
+
+            audio.seek(int(start * original_sr))
+            frames_to_read = min(max_frames, len(audio))
+            data = audio.read(frames_to_read)
+
+            if data.max() > 1 or data.min() < -1:
+                data = data / max(abs(data.max()), abs(data.min()))
+        
+        if original_sr != target_sr:
+            if channels == 1:
+                data = librosa_resample(data.flatten(), orig_sr=original_sr, target_sr=target_sr)
+            else:
+                data = librosa_resample(data.T, orig_sr=original_sr, target_sr=target_sr)[0]
+        else:
+            if channels != 1:
+                data = data.T[0]
+    
+    if data.min() >= 0:
+        data = 2 * data / abs(data.max()) - 1.0
+    else:
+        data = data / max(abs(data.max()), abs(data.min()))
+    
+    assert len(data.shape) == 1, data.shape
+    return data
+
+def process_images(images, image_processor, model_cfg, enable_dynamic_res=False, max_tiles=None):
+    model_cfg.image_processor = image_processor
+    new_images = [
+        process_image(image, model_cfg, None, enable_dynamic_res=enable_dynamic_res, max_tiles=max_tiles)
+        for image in images
+    ]
+
+    if all(x.shape == new_images[0].shape for x in new_images):
+        if len(new_images[0].shape) == 4:
+            new_images = torch.cat(new_images, dim=0)
+        elif len(new_images[0].shape) == 3:
+            new_images = torch.stack(new_images, dim=0)
+        else:
+            raise ValueError(f"new_images rank does not equal to 4, rank: {len(new_images[0].shape)}")
+    else:
+        raise ValueError("The shape of images in new_images is different!")
+    return new_images
+
+def process_sounds(sounds):
+    sounds = torch.tensor(sounds)
+    return sounds
+
+def process_sound_masks(masks):
+    masks = torch.tensor(masks[0])
+    return masks
+
+def tokenizer_image_token(prompt, tokenizer, return_tensors=None):
+    return tokenizer(prompt, return_tensors=return_tensors).input_ids[0]
+
+def is_gemma_tokenizer(tokenizer):
+    return "gemma" in tokenizer.__class__.__name__.lower()
+
+
+def get_model_name_from_path(model_path):
+    model_path = model_path.strip("/")
+    model_paths = model_path.split("/")
+    if model_paths[-1].startswith("checkpoint-"):
+        return model_paths[-2] + "_" + model_paths[-1]
+    else:
+        return model_paths[-1]
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.keyword_ids = []
+        self.max_keyword_len = 0
+        for keyword in keywords:
+            cur_keyword_ids = tokenizer(keyword).input_ids
+            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+                cur_keyword_ids = cur_keyword_ids[1:]
+            if len(cur_keyword_ids) > self.max_keyword_len:
+                self.max_keyword_len = len(cur_keyword_ids)
+            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+        self.tokenizer = tokenizer
+        self.start_len = input_ids.shape[1]
+
+    def call_for_batch(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+        for keyword_id in self.keyword_ids:
+            if (output_ids[0, -keyword_id.shape[0] :] == keyword_id).all():
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        for keyword in self.keywords:
+            if keyword in outputs:
+                return True
+        return False
+
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        outputs = []
+        for i in range(output_ids.shape[0]):
+            outputs.append(self.call_for_batch(output_ids[i].unsqueeze(0), scores))
+        return all(outputs)

llava/model/FloatPointQuantizeTorch.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import math
+
+import torch
+
+
+def floatExMy_quantize_torch(x, e_bit, m_bit, stochastic):
+    sign, x_abs = x.sign(), x.abs()
+    Elow, Ehigh, Mhigh = -(2 ** (e_bit - 1)) + 2, 2 ** (e_bit - 1), 2**m_bit
+    expo = torch.floor(torch.log2(x_abs))
+    expo = torch.clamp(expo, min=Elow, max=Ehigh)
+    mant = x_abs / torch.exp2(expo)
+
+    mant_int = torch.floor(mant)
+    mant_frac = mant - mant_int
+    mant_frac = mant_frac * Mhigh
+    if stochastic:
+        noise = mant_frac.new(mant_frac.shape).uniform_(-0.5, 0.5)
+        mant_frac.add_(noise)
+    mant_frac = torch.round(mant_frac)
+
+    mant_q = mant_int + mant_frac / Mhigh
+    y = sign * (2**expo) * mant_q
+    y = y.to(x)
+
+    return y
+
+
+def floatExM0_quantize_torch(x, e_bit, stochastic):
+    sign, x_abs = x.sign(), x.abs()
+    Elow, Ehigh = -(2 ** (e_bit - 1)) + 1, 2 ** (e_bit - 1)
+    expo = torch.log2(x_abs)
+    if stochastic:
+        noise = expo.new(expo.shape).uniform_(-0.5, 0.5)
+        expo.add(noise)
+        log_bias = math.log2(4 / 3) - 1 / 2
+        expo.add(torch.ones_like(expo) * log_bias)
+    expo = torch.clamp(expo, min=Elow - 1, max=Ehigh)
+    expo = torch.round(expo)
+
+    y = sign * (2**expo) * (expo > Elow)  # When underflow, set the value to 0
+    y = y.to(x)
+
+    return y
+
+
+def Dynamic_quantize_torch(x, bit, stochastic):
+    if stochastic:
+        raise NotImplementedError("Dynamic Tree quantization does not support stochastic")
+    sign, x_abs = x.sign(), x.abs()
+    expo = torch.ceil(torch.log10(x_abs))
+    expo = torch.clamp(expo, min=2 - bit)
+    mant = (10 * x_abs / torch.pow(10, expo) - 1) / 9  # Range from 0 - 1
+
+    mant_frac = mant * 2 ** (bit - 2 - expo.abs())
+    mant_frac = torch.round(mant_frac)
+    mant_frac = mant_frac / (2 ** (bit - 2 - expo.abs())) * 9 + 1
+    y = sign * (10**expo) * mant_frac / 10
+
+    zero_mask = y.abs() > 1.01 * 10 ** (1 - bit)
+    y = y * zero_mask
+    y = y.to(x)
+    return y
+
+
+def ZeroDynamic_quantize_torch(x, bit, stochastic):
+    if stochastic:
+        raise NotImplementedError("Dynamic Tree quantization does not support stochastic")
+    sign, x_abs = x.sign(), x.abs()
+    expo = torch.ceil(torch.log10(x_abs))
+    expo = torch.clamp(expo, min=2 - bit)
+    mant = (10 * x_abs / torch.pow(10, expo) - 1) / 9  # Range from 0 - 1
+
+    mant_frac = mant * 2 ** (bit - 2 - expo.abs())
+    mant_frac = torch.round(mant_frac)
+    mant_frac = mant_frac / (2 ** (bit - 2 - expo.abs())) * 9 + 1
+    y = sign * (10**expo) * mant_frac / 10
+
+    y = y.to(x)
+    return y

llava/model/FloatPointQuantizeTriton.py

@@ -0,0 +1,199 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import math
+import struct
+
+import numpy as np
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+segment_size = 1024**3
+
+
+def floatExMy_quantize_triton(x, e_bit, m_bit, stochastic):
+    x_ori_shape = x.shape
+    x = x.view(-1)
+
+    n_elements = x.numel()
+
+    if n_elements <= segment_size:
+        grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+        y = torch.empty_like(x)
+
+        if x.dtype in [torch.bfloat16, torch.float32]:
+            if stochastic:
+                noise = x.new(x.shape).uniform_(-0.5, 0.5)
+                _floatExMy_stochastic_quantize_kernel[grid](x, noise, y, n_elements, e_bit, m_bit)
+            else:
+                _floatExMy_quantize_kernel[grid](x, y, n_elements, e_bit, m_bit)
+                torch.cuda.synchronize()
+        else:
+            raise NotImplementedError(f"Other data format {x.dtype} for float quantization triton")
+    else:  # Triton will break when x.numel > 2 * 1024 ** 3
+        num_segments = n_elements // segment_size + 1
+        split_size = [segment_size] * (num_segments - 1) + [n_elements - segment_size * (num_segments - 1)]
+        x_list = x.split(split_size)
+        y_list = []
+        del x
+
+        for x in x_list:
+            n_elements = x.numel()
+            grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+            y = torch.empty_like(x)
+
+            if x.dtype in [torch.bfloat16, torch.float32]:
+                if stochastic:
+                    noise = x.new(x.shape).uniform_(-0.5, 0.5)
+                    _floatExMy_stochastic_quantize_kernel[grid](x, noise, y, n_elements, e_bit, m_bit)
+                else:
+                    _floatExMy_quantize_kernel[grid](x, y, n_elements, e_bit, m_bit)
+                    torch.cuda.synchronize()
+            else:
+                raise NotImplementedError(f"Other data format {x.dtype} for float quantization triton")
+
+            y_list.append(y)
+        y = torch.concat(y_list)
+        del y_list
+
+    y = y.reshape(x_ori_shape)
+    return y
+
+
+@triton.autotune(
+    configs=[
+        # triton.Config({'BLOCK_SIZE': 4,}, num_warps=4),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 1024,
+            },
+            num_warps=4,
+        ),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 2048,
+            },
+            num_warps=4,
+        ),
+    ],
+    key=["n_elements"],
+)
+@triton.jit
+def _floatExMy_quantize_kernel(
+    x_ptr,
+    output_ptr,
+    n_elements,
+    e_bit,
+    m_bit,
+    BLOCK_SIZE: tl.constexpr,
+):
+    if isinstance(e_bit, tl.constexpr):
+        ebit = e_bit.value
+    else:
+        ebit = e_bit
+
+    if isinstance(m_bit, tl.constexpr):
+        mbit = m_bit.value
+    else:
+        mbit = m_bit
+
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+
+    x = x.to(tl.float32)
+    sign = 1 - 2 * libdevice.signbit(x)
+    x_abs = tl.abs(x)
+    Elow = -tl.exp2((ebit - 1).to(tl.float32)) + 2
+    Ehigh = tl.exp2((ebit - 1).to(tl.float32))
+    Mhigh = tl.exp2(mbit.to(tl.float32))
+    expo = tl.floor(tl.log2(x_abs))
+    expo = tl.clamp(expo, min=Elow, max=Ehigh)
+    mant = x_abs / tl.exp2(expo)
+
+    mant_int = tl.floor(mant)
+    mant_frac = mant - mant_int
+    mant_frac = mant_frac * Mhigh
+    # mant_frac = mant_frac + noise
+    mant_frac = libdevice.round(mant_frac)
+
+    mant_q = mant_int + mant_frac / Mhigh
+    y = sign * tl.exp2(expo) * mant_q
+    y = y.to(x_ptr.dtype.element_ty)
+
+    tl.store(output_ptr + offsets, y, mask=mask)
+
+
+@triton.autotune(
+    configs=[
+        # triton.Config({'BLOCK_SIZE': 4,}, num_warps=4),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 1024,
+            },
+            num_warps=4,
+        ),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 2048,
+            },
+            num_warps=4,
+        ),
+    ],
+    key=["n_elements"],
+)
+@triton.jit
+def _floatExMy_stochastic_quantize_kernel(
+    x_ptr,
+    noise_ptr,
+    output_ptr,
+    n_elements,
+    e_bit,
+    m_bit,
+    BLOCK_SIZE: tl.constexpr,
+):
+    if isinstance(e_bit, tl.constexpr):
+        ebit = e_bit.value
+    else:
+        ebit = e_bit
+
+    if isinstance(m_bit, tl.constexpr):
+        mbit = m_bit.value
+    else:
+        mbit = m_bit
+
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+    noise = tl.load(noise_ptr + offsets, mask=mask)
+
+    x = x.to(tl.float32)
+    sign = 1 - 2 * libdevice.signbit(x)
+    x_abs = tl.abs(x)
+    Elow = -tl.exp2((ebit - 1).to(tl.float32)) + 2
+    Ehigh = tl.exp2((ebit - 1).to(tl.float32))
+    Mhigh = tl.exp2(mbit.to(tl.float32))
+    expo = tl.floor(tl.log2(x_abs))
+    expo = tl.clamp(expo, min=Elow, max=Ehigh)
+    mant = x_abs / tl.exp2(expo)
+
+    mant_int = tl.floor(mant)
+    mant_frac = mant - mant_int
+    mant_frac = mant_frac * Mhigh
+    mant_frac = mant_frac + noise
+    mant_frac = libdevice.round(mant_frac)
+
+    mant_q = mant_int + mant_frac / Mhigh
+    y = sign * tl.exp2(expo) * mant_q
+    y = y.to(x_ptr.dtype.element_ty)
+
+    tl.store(output_ptr + offsets, y, mask=mask)

llava/model/__init__.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+from .language_model.llava_llama import LlavaLlamaConfig, LlavaLlamaModel
+
+# FP8 related comments, development in progress (PI: ligeng zhu, haochen xi)
+# NOTE: VLM + LLM
+# from .language_model.qllava_qllama import QLlavaLlamaConfig, QLlavaLlamaModel
+# NOTE: Linear -> fp8, similar to transformer engine
+# from .language_model.qllama import QLlamaConfig, QLlamaForCausalLM, QLlamaModel
+# NOTE: Linear + Activation -> fp8, haochen's iclr version
+# from .language_model.qmemllama import QMemLlamaConfig, QMemLlamaForCausalLM, QMemLlamaModel
+"""
+TODO:
+    linear(weights):
+        simulated fp8: done
+        real fp8: in-progress (code already implmented)
+    activation:
+        simulated fp8: done
+        real fp8: in-progress (still coding)
+    optimizers:
+        current VILA: bf16
+        simulated fp8: done
+        real fp8 + fsdp (single node): done
+        real fp8 + fsdp (multiple node): in-progress
+1. linear fp8
+2. activation fp8
+3. fp8 infernce example (load directly from a fp8 and fwd)
+4. bind fp8 related configs to QLlamaConfig {"coat_fp8_args": {}}
+"""
+from .language_model.fp8linearqwen2 import FP8LinearQwen2Config, FP8LinearQwen2Model
+from .language_model.qllava_qllama import QLlavaLlamaConfig, QLlavaLlamaModel

llava/model/apply_delta.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# This file is modified from https://github.com/haotian-liu/LLaVA/
+
+"""
+Usage:
+python3 -m fastchat.model.apply_delta --base ~/model_weights/llama-7b --target ~/model_weights/vicuna-7b --delta lmsys/vicuna-7b-delta
+"""
+import argparse
+
+import torch
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from llava import LlavaLlamaForCausalLM
+
+
+def apply_delta(base_model_path, target_model_path, delta_path):
+    print("Loading base model")
+    base = AutoModelForCausalLM.from_pretrained(base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+
+    print("Loading delta")
+    delta = LlavaLlamaForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)
+
+    print("Applying delta")
+    for name, param in tqdm(delta.state_dict().items(), desc="Applying delta"):
+        if name not in base.state_dict():
+            assert name in [
+                "model.mm_projector.weight",
+                "model.mm_projector.bias",
+            ], f"{name} not in base model"
+            continue
+        if param.data.shape == base.state_dict()[name].shape:
+            param.data += base.state_dict()[name]
+        else:
+            assert name in [
+                "model.embed_tokens.weight",
+                "lm_head.weight",
+            ], f"{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}"
+            bparam = base.state_dict()[name]
+            param.data[: bparam.shape[0], : bparam.shape[1]] += bparam
+
+    print("Saving target model")
+    delta.save_pretrained(target_model_path)
+    delta_tokenizer.save_pretrained(target_model_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-model-path", type=str, required=True)
+    parser.add_argument("--target-model-path", type=str, required=True)
+    parser.add_argument("--delta-path", type=str, required=True)
+
+    args = parser.parse_args()
+
+    apply_delta(args.base_model_path, args.target_model_path, args.delta_path)

llava/model/builder.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# This file is modified from https://github.com/haotian-liu/LLaVA/
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+import os
+import warnings
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, PretrainedConfig
+
+from llava.model import LlavaLlamaModel
+from llava.model.utils import is_mm_model
+
+
+def load_pretrained_model(
+    model_path,
+    model_name,
+    model_base=None,
+    load_8bit=False,
+    load_4bit=False,
+    device_map="auto",
+    device="cuda",
+    **kwargs,
+):
+    kwargs = {"device_map": device_map, **kwargs}
+
+    if device != "cuda":
+        kwargs["device_map"] = {"": device}
+
+    if load_8bit:
+        kwargs["load_in_8bit"] = True
+    elif load_4bit:
+        kwargs["load_in_4bit"] = True
+        kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+    else:
+        kwargs["torch_dtype"] = torch.float16
+        # kwargs["torch_dtype"] = torch.bfloat16
+
+    if is_mm_model(model_path):
+        # Load LLaVA model
+        ## TODO @yunhao: mind fixing lora
+        if "lora" in model_name.lower() and model_base is None:
+            warnings.warn(
+                "There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged."
+            )
+        if ("lora" in model_name.lower() or "dora" in model_name.lower()) and model_base is not None:
+            lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)
+            print(lora_cfg_pretrained)
+            print("Loading LLaVA from base model...")
+            config = AutoConfig.from_pretrained(model_base)
+            prepare_config_for_eval(config, kwargs)
+            model = LlavaLlamaModel.from_pretrained(model_base, low_cpu_mem_usage=True, config=config, **kwargs)
+            tokenizer = model.tokenizer
+            token_num, tokem_dim = model.llm.lm_head.out_features, model.llm.lm_head.in_features
+            if model.llm.lm_head.weight.shape[0] != token_num:
+                model.llm.lm_head.weight = torch.nn.Parameter(
+                    torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype)
+                )
+                model.llm.embed_tokens.weight = torch.nn.Parameter(
+                    torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype)
+                )
+
+            print("Loading additional LLaVA weights...")
+            if os.path.exists(os.path.join(model_path, "non_lora_trainables.bin")):
+                non_lora_trainables = torch.load(
+                    os.path.join(model_path, "non_lora_trainables.bin"),
+                    map_location="cpu",
+                )
+            else:
+                # this is probably from HF Hub
+                from huggingface_hub import hf_hub_download
+
+                def load_from_hf(repo_id, filename, subfolder=None):
+                    cache_file = hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder)
+                    return torch.load(cache_file, map_location="cpu")
+
+                non_lora_trainables = load_from_hf(model_path, "non_lora_trainables.bin")
+            non_lora_trainables = {
+                (k[11:] if k.startswith("base_model.") else k): v for k, v in non_lora_trainables.items()
+            }
+            if any(k.startswith("model.model.") for k in non_lora_trainables):
+                non_lora_trainables = {
+                    (k[6:] if k.startswith("model.") else k): v for k, v in non_lora_trainables.items()
+                }
+            model.load_state_dict(non_lora_trainables, strict=False)
+
+            from peft import PeftModel
+
+            print("Loading LoRA weights...")
+            model = PeftModel.from_pretrained(model, model_path)
+            print("Merging LoRA weights...")
+            model = model.merge_and_unload()
+            print("Model is loaded...")
+        else:
+            config = AutoConfig.from_pretrained(model_path)
+            config.resume_path = model_path
+            prepare_config_for_eval(config, kwargs)
+            model = LlavaLlamaModel(config=config, low_cpu_mem_usage=True, **kwargs)
+            tokenizer = model.tokenizer
+    else:
+        # Load language model
+        if model_base is not None:
+            # PEFT model
+            from peft import PeftModel
+
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, **kwargs)
+            print(f"Loading LoRA weights from {model_path}")
+            model = PeftModel.from_pretrained(model, model_path)
+            print(f"Merging weights")
+            model = model.merge_and_unload()
+            print("Convert to FP16...")
+            model.to(torch.float16)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, legacy=False)
+            model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+    model.eval()
+    image_processor = None
+    if is_mm_model(model_path):
+        model.resize_token_embeddings(len(tokenizer))
+
+    if hasattr(model.llm.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    else:
+        context_len = 2048
+
+    return tokenizer, model, image_processor, context_len
+
+
+def prepare_config_for_eval(config: PretrainedConfig, kwargs: dict):
+    try:
+        # compatible with deprecated config convention
+        if getattr(config, "vision_tower_cfg", None) is None:
+            config.vision_tower_cfg = config.mm_vision_tower
+    except AttributeError:
+        raise ValueError(f"Invalid configuration! Cannot find vision_tower in config:\n{config}")
+
+    config.model_dtype = kwargs.pop("torch_dtype").__str__()

llava/model/coat/activation/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+

llava/model/coat/activation/fake_quantization/FloatPointQuantizeTorch.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+
+import torch
+
+
+def floatExMy_quantize_torch(x, e_bit, m_bit, stochastic):
+    sign, x_abs = x.sign(), x.abs()
+    Elow, Ehigh, Mhigh = -(2 ** (e_bit - 1)) + 2, 2 ** (e_bit - 1), 2**m_bit
+    expo = torch.floor(torch.log2(x_abs))
+    expo = torch.clamp(expo, min=Elow, max=Ehigh)
+    mant = x_abs / torch.exp2(expo)
+
+    mant_int = torch.floor(mant)
+    mant_frac = mant - mant_int
+    mant_frac = mant_frac * Mhigh
+    if stochastic:
+        noise = mant_frac.new(mant_frac.shape).uniform_(-0.5, 0.5)
+        mant_frac.add_(noise)
+    mant_frac = torch.round(mant_frac)
+
+    mant_q = mant_int + mant_frac / Mhigh
+    y = sign * (2**expo) * mant_q
+    y = y.to(x)
+
+    return y
+
+
+def floatExM0_quantize_torch(x, e_bit, stochastic):
+    sign, x_abs = x.sign(), x.abs()
+    Elow, Ehigh = -(2 ** (e_bit - 1)) + 1, 2 ** (e_bit - 1)
+    expo = torch.log2(x_abs)
+    if stochastic:
+        noise = expo.new(expo.shape).uniform_(-0.5, 0.5)
+        expo.add(noise)
+        log_bias = math.log2(4 / 3) - 1 / 2
+        expo.add(torch.ones_like(expo) * log_bias)
+    expo = torch.clamp(expo, min=Elow - 1, max=Ehigh)
+    expo = torch.round(expo)
+
+    y = sign * (2**expo) * (expo > Elow)  # When underflow, set the value to 0
+    y = y.to(x)
+
+    return y
+
+
+def Dynamic_quantize_torch(x, bit, stochastic):
+    if stochastic:
+        raise NotImplementedError("Dynamic Tree quantization does not support stochastic")
+    sign, x_abs = x.sign(), x.abs()
+    expo = torch.ceil(torch.log10(x_abs))
+    expo = torch.clamp(expo, min=2 - bit)
+    mant = (10 * x_abs / torch.pow(10, expo) - 1) / 9  # Range from 0 - 1
+
+    mant_frac = mant * 2 ** (bit - 2 - expo.abs())
+    mant_frac = torch.round(mant_frac)
+    mant_frac = mant_frac / (2 ** (bit - 2 - expo.abs())) * 9 + 1
+    y = sign * (10**expo) * mant_frac / 10
+
+    zero_mask = y.abs() > 1.01 * 10 ** (1 - bit)
+    y = y * zero_mask
+    y = y.to(x)
+    return y
+
+
+def ZeroDynamic_quantize_torch(x, bit, stochastic):
+    if stochastic:
+        raise NotImplementedError("Dynamic Tree quantization does not support stochastic")
+    sign, x_abs = x.sign(), x.abs()
+    expo = torch.ceil(torch.log10(x_abs))
+    expo = torch.clamp(expo, min=2 - bit)
+    mant = (10 * x_abs / torch.pow(10, expo) - 1) / 9  # Range from 0 - 1
+
+    mant_frac = mant * 2 ** (bit - 2 - expo.abs())
+    mant_frac = torch.round(mant_frac)
+    mant_frac = mant_frac / (2 ** (bit - 2 - expo.abs())) * 9 + 1
+    y = sign * (10**expo) * mant_frac / 10
+
+    y = y.to(x)
+    return y

llava/model/coat/activation/fake_quantization/FloatPointQuantizeTriton.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import struct
+
+import numpy as np
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+
+def floatExMy_quantize_triton(x, e_bit, m_bit, stochastic):
+    n_elements = x.numel()
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+    y = torch.zeros_like(x)
+
+    if x.dtype in [torch.bfloat16, torch.float32]:
+        if stochastic:
+            noise = x.new(x.shape).uniform_(-0.5, 0.5)
+            _floatExMy_stochastic_quantize_kernel[grid](x, noise, y, n_elements, e_bit, m_bit)
+        else:
+            _floatExMy_quantize_kernel[grid](x, y, n_elements, e_bit, m_bit)
+    else:
+        raise NotImplementedError(f"Other data format {x.dtype} for float quantization triton")
+
+    return y
+
+
+@triton.autotune(
+    configs=[
+        # triton.Config({'BLOCK_SIZE': 4,}, num_warps=4),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 1024,
+            },
+            num_warps=4,
+        ),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 2048,
+            },
+            num_stages=1,
+        ),
+    ],
+    key=["n_elements"],
+)
+@triton.jit
+def _floatExMy_quantize_kernel(
+    x_ptr,
+    output_ptr,
+    n_elements,
+    e_bit,
+    m_bit,
+    BLOCK_SIZE: tl.constexpr,
+):
+    if isinstance(e_bit, tl.constexpr):
+        ebit = e_bit.value
+    else:
+        ebit = e_bit
+
+    if isinstance(m_bit, tl.constexpr):
+        mbit = m_bit.value
+    else:
+        mbit = m_bit
+
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+
+    x = x.to(tl.float32)
+    sign = 1 - 2 * libdevice.signbit(x)
+    x_abs = tl.abs(x)
+    Elow = -tl.exp2((ebit - 1).to(tl.float32)) + 2
+    Ehigh = tl.exp2((ebit - 1).to(tl.float32))
+    Mhigh = tl.exp2(mbit.to(tl.float32))
+    expo = tl.floor(tl.log2(x_abs))
+    expo = tl.clamp(expo, min=Elow, max=Ehigh)
+    mant = x_abs / tl.exp2(expo)
+
+    mant_int = tl.floor(mant)
+    mant_frac = mant - mant_int
+    mant_frac = mant_frac * Mhigh
+    # mant_frac = mant_frac + noise
+    mant_frac = libdevice.round(mant_frac)
+
+    mant_q = mant_int + mant_frac / Mhigh
+    y = sign * tl.exp2(expo) * mant_q
+    y = y.to(x_ptr.dtype.element_ty)
+
+    tl.store(output_ptr + offsets, y, mask=mask)
+
+
+@triton.autotune(
+    configs=[
+        # triton.Config({'BLOCK_SIZE': 4,}, num_warps=4),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 1024,
+            },
+            num_warps=4,
+        ),
+        triton.Config(
+            {
+                "BLOCK_SIZE": 2048,
+            },
+            num_stages=1,
+        ),
+    ],
+    key=["n_elements"],
+)
+@triton.jit
+def _floatExMy_stochastic_quantize_kernel(
+    x_ptr,
+    noise_ptr,
+    output_ptr,
+    n_elements,
+    e_bit,
+    m_bit,
+    BLOCK_SIZE: tl.constexpr,
+):
+    if isinstance(e_bit, tl.constexpr):
+        ebit = e_bit.value
+    else:
+        ebit = e_bit
+
+    if isinstance(m_bit, tl.constexpr):
+        mbit = m_bit.value
+    else:
+        mbit = m_bit
+
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+    noise = tl.load(noise_ptr + offsets, mask=mask)
+
+    x = x.to(tl.float32)
+    sign = 1 - 2 * libdevice.signbit(x)
+    x_abs = tl.abs(x)
+    Elow = -tl.exp2((ebit - 1).to(tl.float32)) + 2
+    Ehigh = tl.exp2((ebit - 1).to(tl.float32))
+    Mhigh = tl.exp2(mbit.to(tl.float32))
+    expo = tl.floor(tl.log2(x_abs))
+    expo = tl.clamp(expo, min=Elow, max=Ehigh)
+    mant = x_abs / tl.exp2(expo)
+
+    mant_int = tl.floor(mant)
+    mant_frac = mant - mant_int
+    mant_frac = mant_frac * Mhigh
+    mant_frac = mant_frac + noise
+    mant_frac = libdevice.round(mant_frac)
+
+    mant_q = mant_int + mant_frac / Mhigh
+    y = sign * tl.exp2(expo) * mant_q
+    y = y.to(x_ptr.dtype.element_ty)
+
+    tl.store(output_ptr + offsets, y, mask=mask)

llava/model/coat/activation/fake_quantization/quantize_function.py

@@ -0,0 +1,239 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+import torch
+
+from .FloatPointQuantizeTorch import *
+from .FloatPointQuantizeTriton import *
+
+
+def block_cut(input, row_block, column_block, pad_block=False):
+    # print(input.shape)
+    original_shape = input.shape
+    # input tensor shape is M * N
+    if len(input.shape) > 2:
+        input = input.reshape(-1, input.shape[2])
+    elif len(input.shape) == 2:
+        pass
+    else:
+        raise ValueError(f"input shape {input.shape} does not match for block cut, {input}")
+    M, N = input.shape[0], input.shape[1]
+
+    if row_block == -1:
+        row_block = M
+    if column_block == -1:
+        column_block = N
+
+    if pad_block:
+        row_remainder, col_remainder = M % row_block, N % column_block
+        if row_remainder:
+            row_pad = row_block - row_remainder
+        else:
+            row_pad = 0
+        if col_remainder:
+            col_pad = column_block - col_remainder
+        else:
+            col_pad = 0
+
+        input = torch.nn.functional.pad(
+            input, (0, col_pad, 0, row_pad), "constant", 0
+        )  # refer to torch's doc to see why
+        M, N = input.shape[0], input.shape[1]
+        row_num, column_num = M // row_block, N // column_block
+    else:
+        row_num, column_num = M // row_block, N // column_block
+
+    assert row_num * row_block == M, f"{row_num}, {row_block}, {M}, {original_shape}"
+    assert column_num * column_block == N, f"{column_num}, {column_block}, {N}, {original_shape}"
+    # print(input.shape)
+    input = (
+        input.reshape(row_num, row_block, column_num, column_block)
+        .permute(0, 2, 1, 3)
+        .reshape(row_num * column_num, row_block, column_block)
+    )
+    # print(input.shape)
+    return input
+
+
+def block_reshape(input, origin_input, row_block, column_block, pad_block=False):
+    if len(origin_input.shape) > 2:
+        flatten_input = origin_input.reshape(-1, origin_input.shape[2])
+    elif len(origin_input.shape) == 2:
+        flatten_input = origin_input
+    else:
+        raise ValueError(f"input shape {input.shape} does not match for block cut")
+
+    M, N = flatten_input.shape[0], flatten_input.shape[1]
+
+    if row_block == -1:
+        row_block = M
+    if column_block == -1:
+        column_block = N
+
+    if pad_block:
+        row_remainder, col_remainder = M % row_block, N % column_block
+        if row_remainder:
+            row_pad = row_block - row_remainder
+        else:
+            row_pad = 0
+        if col_remainder:
+            col_pad = column_block - col_remainder
+        else:
+            col_pad = 0
+
+        pad_origin_input = torch.nn.functional.pad(origin_input, (0, col_pad, 0, row_pad), "constant", 0)
+        M, N = pad_origin_input.shape[0], pad_origin_input.shape[1]
+        row_num, column_num = M // row_block, N // column_block
+    else:
+        row_num, column_num = M // row_block, N // column_block
+
+    input = (
+        input.reshape(row_num, column_num, row_block, column_block)
+        .permute(0, 2, 1, 3)
+        .reshape(row_num * row_block, column_num * column_block)
+    )
+
+    M, N = flatten_input.shape[0], flatten_input.shape[1]
+    input = input[:M, :N]
+
+    if len(origin_input.shape) > 2:
+        input = input.reshape(origin_input.shape)
+    elif len(origin_input.shape) == 2:
+        pass
+    else:
+        raise ValueError(f"input shape {input.shape} does not match for block reshape")
+
+    return input
+
+
+def block_verify_int8(input, row_block, column_block, layer_type, necessary=True):
+    Binput = block_cut(input, row_block, column_block)
+    Binput = Binput.to(torch.float32)
+
+    for n in range(Binput.shape[0]):
+        unique_values = len(torch.unique(Binput[n, :, :]))
+        if unique_values > 256:
+            if necessary:
+                raise ValueError(f"{layer_type} contains more than 256 unique values.")
+            else:
+                return False
+    return True
+
+
+def block_quant(input, symm, bits, stochastic, epsilon, apply_quantize, layer_name):
+    Quant_fn = SymmQuantizer
+    return Quant_fn.apply(input, symm, bits, stochastic, epsilon, apply_quantize, layer_name)
+
+
+def extract_bit(string):
+    match = re.match(r"INT(\d+)", string)  # INT8
+    if match:
+        return "integer", int(match.group(1)), None
+    match = re.match(r"E(\d+)M(\d+)", string)  # E4M3 / E5M2
+    if match:
+        Ebit, Mbit = int(match.group(1)), int(match.group(2))
+        if Ebit == 1:
+            return "integer", Mbit + 1, None
+        if Mbit == 0:
+            return "floatExM0", int(match.group(1)), 0
+        return "floatExMy", int(match.group(1)), int(match.group(2))
+    match = re.match(r"DE(\d+)", string)
+    if match:
+        return "Dynamic", int(match.group(1)), None
+    match = re.match(r"ZeroD(\d+)", string)
+    if match:
+        return "ZeroDynamic", int(match.group(1)), None
+    raise ValueError(f"{string} data format is not supported")
+
+
+class SymmQuantizer(torch.autograd.function.InplaceFunction):
+    @staticmethod
+    def forward(ctx, input, symm, bits, stochastic, epsilon, apply_quantize=True, layer_name=None):
+        with torch.no_grad():
+            absmax_per_block = input.abs().amax(dim=(1, 2)).unsqueeze(1).unsqueeze(2) + epsilon
+
+            if bits == "100" or not apply_quantize:
+                return input, input, torch.ones_like(absmax_per_block)
+            elif bits == "FP32":
+                return input.to(torch.float32), input.to(torch.float32), torch.ones_like(absmax_per_block)
+            elif bits == "FP16":
+                return input.to(torch.float16), input.to(torch.float16), torch.ones_like(absmax_per_block)
+            elif bits == "BF16":
+                return input.to(torch.bfloat16), input.to(torch.bfloat16), torch.ones_like(absmax_per_block)
+            else:
+                QuantType, bit1, bit2 = extract_bit(bits)
+                if not symm:
+                    bit1 = bit1 + 1  # pretend to be asymmtric
+
+                if QuantType == "integer":
+                    Qn, Qp = -(2 ** (bit1 - 1) - 1), 2 ** (bit1 - 1) - 1
+                elif QuantType == "floatExMy":
+                    Qn, Qp = -(2 - 2 ** (-bit2)) * (2 ** (2 ** (bit1 - 1))), (2 - 2 ** (-bit2)) * (
+                        2 ** (2 ** (bit1 - 1))
+                    )
+                    if bit1 == 4 and bit2 == 3:  # E4M3
+                        Qn, Qp = -448, 448
+                    if bit1 == 5 and bit2 == 2:  # E5M2
+                        Qn, Qp = -57344, 57344
+                elif QuantType == "floatExM0":
+                    Qn, Qp = -(2 ** (2 ** (bit1 - 1))) + 1, 2 ** (2 ** (bit1 - 1))
+                elif QuantType == "Dynamic":
+                    Qn, Qp = -1, 1
+                elif QuantType == "ZeroDynamic":
+                    Qn, Qp = -1, 1
+                else:
+                    raise NotImplementedError(f"{bits} is not supported by quantization")
+                scale_per_block = (2 * absmax_per_block) / (Qp - Qn)
+                scale_per_block = scale_per_block.to(input)
+
+                Qinput = input / scale_per_block
+
+                if QuantType == "integer":
+                    if stochastic:
+                        noise = Qinput.new(Qinput.shape).uniform_(-0.5, 0.5)
+                        Qinput.add_(noise)
+                    Qinput.clamp_(Qn, Qp).round_()
+                elif QuantType == "floatExMy":
+                    # Qinput = floatExMy_quantize_torch(Qinput, bit1, bit2, stochastic)
+                    Qinput = floatExMy_quantize_triton(Qinput, bit1, bit2, stochastic)
+                elif QuantType == "floatExM0":
+                    Qinput = floatExM0_quantize_torch(Qinput, bit1, stochastic)
+                else:
+                    raise NotImplementedError(f"{bits} is not supported by quantization")
+
+                RQinput = Qinput * scale_per_block
+
+                if input.dtype != Qinput.dtype:
+                    print(
+                        f"Input type is {input.dtype}, Qinput type is {Qinput.dtype}, scale_per_block type is {scale_per_block.dtype}",
+                        file=open("debug.txt", "a"),
+                    )
+                    import IPython
+
+                    IPython.embed()
+                return RQinput, Qinput, scale_per_block
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output, None, None, None, None, None

llava/model/coat/activation/fake_quantization/utils.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+def list_has_common_element(list1, list2):
+    set1 = set(list1)
+    set2 = set(list2)
+    return len(set1.intersection(set2)) > 0
+
+
+def calculate_scale_num(input, row_block, col_block):
+    if len(input.shape) > 2:
+        input = input.reshape(-1, input.shape[2])
+    elif len(input.shape) == 2:
+        pass
+    else:
+        raise ValueError(f"input shape {input.shape} does not match for block cut, {input}")
+    M, N = input.shape[0], input.shape[1]
+
+    if row_block == -1:
+        row_block = M
+    if col_block == -1:
+        col_block = N
+
+    return input.numel() / (row_block * col_block)
+
+
+def quant_get_local_rank() -> int:
+    return int(os.environ.get("LOCAL_RANK") or 0)
+
+
+def format_string_with_condition(
+    input_string,
+    condition_config,
+    symm,
+    bits,
+    blocksize_config,
+    input_pad=20,
+):
+    padded_string = input_string.ljust(input_pad)
+    output_string = padded_string
+
+    for k, v in condition_config.items():
+        if v:
+            output_string = output_string + k.ljust(10) + "True".ljust(6) + "".ljust(6)
+        else:
+            output_string = output_string + k.ljust(10) + "".ljust(6) + "False".ljust(6)
+
+    output_string = output_string + f"Symm {symm}".ljust(10)
+
+    for k, v in bits.items():
+        output_string = output_string + f"{k} bit".ljust(10) + v.ljust(10)
+    for k, v in blocksize_config.items():
+        output_string += f"{k}: {v}".ljust(15)
+
+    return output_string
+
+
+def print_warning(sentence):
+    print("*" * (len(sentence) + 4))
+    print(f"* {sentence} *")
+    print("*" * (len(sentence) + 4))
+
+
+def check_nan_inf(tensor, check_nan, check_inf):
+    if check_nan:
+        contain_nan = torch.isnan(tensor).any()
+    else:
+        contain_nan = False
+    if check_inf:
+        contain_inf = torch.isinf(tensor).any()
+    else:
+        contain_inf = False
+    return contain_nan, contain_inf
+
+
+def move_torch_to_numpy(tensor):
+    if tensor is None:
+        return None
+
+    if tensor.is_cuda:
+        tensor = tensor.cpu()
+    return tensor.detach().float().numpy()
+
+
+def flatten_to_1d(tensor):
+    if tensor is None:
+        return None
+
+    return tensor.reshape(-1)

llava/model/coat/activation/models/_fp8_quantization_config.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+from dataclasses import dataclass
+
+from transformers import PretrainedConfig
+
+
+@dataclass
+class QuantizationConfig:
+    quantize_model: str = "false"
+    symm: bool = True
+    epsilon: float = 1e-10
+    fabit: str = "E4M3"
+    fwbit: str = "E4M3"
+    fobit: str = "E4M3"
+    babit: str = "E5M2"
+    bwbit: str = "E5M2"
+    bobit: str = "E5M2"
+    qchoice: str = "none"
+    group_size: int = -1
+    pad_to_multiple_of: int = 0
+    weight_memory_efficient: bool = True
+
+    # Legacy
+    row_blocksize: int = -1
+    col_blocksize: int = -1
+
+    def __init__(
+        self,
+        quantize_model: str = "false",
+        symm: bool = True,
+        epsilon: float = 1e-10,
+        fabit: str = "E4M3",
+        fwbit: str = "E4M3",
+        fobit: str = "E4M3",
+        babit: str = "E5M2",
+        bwbit: str = "E5M2",
+        bobit: str = "E5M2",
+        qchoice: str = "none",
+        group_size: int = -1,
+        pad_to_multiple_of: int = 0,
+        weight_memory_efficient: bool = True,
+        row_blocksize: int = -1,
+        col_blocksize: int = -1,
+        **kwargs,
+    ):
+        super().__init__()
+        self.quantize_model = quantize_model
+        self.symm = symm
+        self.epsilon = epsilon
+        self.fabit = fabit
+        self.fwbit = fwbit
+        self.fobit = fobit
+        self.babit = babit
+        self.bwbit = bwbit
+        self.bobit = bobit
+        self.qchoice = qchoice
+        self.group_size = group_size
+        self.pad_to_multiple_of = pad_to_multiple_of
+        self.weight_memory_efficient = weight_memory_efficient
+
+        self.row_blocksize = row_blocksize
+        self.col_blocksize = col_blocksize

llava/model/coat/activation/models/_fp8_weightcache.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+
+from ..real_quantization import fp8_division_transpose
+
+
+class FP8CacheWeightModule(nn.Module):
+    def __init__(self, config, qargs, layer_id):
+        super().__init__()
+        self.config = config
+        self.qargs = qargs
+        self.layer_id = layer_id
+
+    def prepare_weight(self, weight, weight_name, is_first_microbatch):
+        if is_first_microbatch:
+            if self.qargs.weight_memory_efficient:
+                # print(f"{weight_name} uses first microbatch")
+                weight_fp8, weight_s, weight_fp8_t = fp8_division_transpose(
+                    weight, self.qargs.group_size, self.fwobits["fwbit"]
+                )
+                setattr(self, f"{weight_name}_fp8_scale", weight_s)
+                return weight_fp8, weight_fp8_t, weight_s
+            else:
+                # print(f"{weight_name} uses first microbatch")
+                weight_fp8, weight_s, weight_fp8_t = fp8_division_transpose(
+                    weight, self.qargs.group_size, self.fwobits["fwbit"]
+                )
+                setattr(self, f"{weight_name}_fp8", weight_fp8)
+                setattr(self, f"{weight_name}_fp8_t", weight_fp8_t)
+                setattr(self, f"{weight_name}_fp8_scale", weight_s)
+                return weight_fp8, weight_fp8_t, weight_s
+        else:
+            if self.qargs.weight_memory_efficient:
+                return getattr(self, f"{weight_name}_fp8_scale")
+            else:
+                return (
+                    getattr(self, f"{weight_name}_fp8"),
+                    getattr(self, f"{weight_name}_fp8_t"),
+                    getattr(self, f"{weight_name}_fp8_scale"),
+                )
+
+    def forward(self, x):
+        pass

llava/model/coat/activation/models/_fp8manager.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+
+class FP8Manager:
+    """Class to keep track of and manipulate the global
+    FP8 state at different stages of execution.
+    """
+
+    is_first_microbatch = False

llava/model/coat/activation/models/coat_llama.py

@@ -0,0 +1,1479 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+import os
+from fnmatch import fnmatch
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache, StaticCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_flash_attention_utils import _flash_attention_forward
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import (
+    LlamaAttention,
+    LlamaDynamicNTKScalingRotaryEmbedding,
+    LlamaForCausalLM,
+    LlamaLinearScalingRotaryEmbedding,
+    LlamaModel,
+    LlamaPreTrainedModel,
+    LlamaRMSNorm,
+    LlamaRotaryEmbedding,
+    _prepare_4d_causal_attention_mask_with_cache_position,
+    apply_rotary_pos_emb,
+    repeat_kv,
+    rotate_half,
+)
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_greater_or_equal_2_10,
+    is_torchdynamo_compiling,
+    logging,
+    replace_return_docstrings,
+)
+
+from ..real_quantization import (
+    Coat_quantize_bgn,
+    Coat_quantize_end,
+    fp8_add_Ifp_Ifp_Ofp_Og16,
+    fp8_add_Ifp_Ifp_Ofp_Opt,
+    fp8_division,
+    fp8_division_transpose,
+    fp8_gelu_backward,
+    fp8_gelu_forward,
+    fp8_layernorm_noparam_backward,
+    fp8_layernorm_noparam_forward,
+    fp8_linear_backward,
+    fp8_linear_forward,
+    fp8_mul_backward,
+    fp8_mul_forward,
+    fp8_quantize,
+    fp8_quantize_pertensor,
+    fp8_quantize_pertensor_transpose,
+    fp8_rmsnorm_backward,
+    fp8_rmsnorm_forward,
+    fp8_silu_backward,
+    fp8_silu_forward,
+    fp8_transpose,
+)
+
+# FP8 related
+from ._fp8_quantization_config import QuantizationConfig
+from ._fp8_weightcache import FP8CacheWeightModule
+from ._fp8manager import FP8Manager
+
+logger = logging.get_logger(__name__)
+
+
+class CoatLlamaConfig(LlamaConfig):
+    model_type = "fp8_llama"
+
+
+class CoatLlamaBeforeAttentionResidual(FP8CacheWeightModule):
+    """
+    This is a typical transformer attention module that contains (1) Residual (2) LayerNorm / RMSNorm (3) 1 * Linear layers
+    """
+
+    def __init__(self, config: CoatLlamaConfig, qargs: QuantizationConfig, layer_idx: Optional[int] = None):
+        super().__init__(config, qargs, layer_idx)
+
+        self.qargs = qargs
+        self.fwobits = {
+            "fabit": self.qargs.fabit,
+            "fwbit": self.qargs.fwbit,
+            "fobit": self.qargs.fobit,
+            "babit": self.qargs.babit,
+            "bwbit": self.qargs.bwbit,
+            "bobit": self.qargs.bobit,
+        }
+
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+
+    def forward(self, re_x, x, s, rmsnorm_weight):
+        if self.training:
+            if self.qargs.weight_memory_efficient:
+                # Prepare
+                with torch.no_grad():
+                    weight1_s = self.prepare_weight(self.q_proj.weight, "q_proj", FP8Manager.is_first_microbatch)
+                    weight2_s = self.prepare_weight(self.k_proj.weight, "k_proj", FP8Manager.is_first_microbatch)
+                    weight3_s = self.prepare_weight(self.v_proj.weight, "v_proj", FP8Manager.is_first_microbatch)
+                return _CoatLlamaBeforeAttentionResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.q_proj.weight,
+                    None,
+                    None,
+                    weight1_s,
+                    self.k_proj.weight,
+                    None,
+                    None,
+                    weight2_s,
+                    self.v_proj.weight,
+                    None,
+                    None,
+                    weight3_s,
+                    rmsnorm_weight,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+            else:
+                # Prepare
+                with torch.no_grad():
+                    weight1, weight1_t, weight1_s = self.prepare_weight(
+                        self.q_proj.weight, "q_proj", FP8Manager.is_first_microbatch
+                    )
+                    weight2, weight2_t, weight2_s = self.prepare_weight(
+                        self.k_proj.weight, "k_proj", FP8Manager.is_first_microbatch
+                    )
+                    weight3, weight3_t, weight3_s = self.prepare_weight(
+                        self.v_proj.weight, "v_proj", FP8Manager.is_first_microbatch
+                    )
+                return _CoatLlamaBeforeAttentionResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.q_proj.weight,
+                    weight1,
+                    weight1_t,
+                    weight1_s,
+                    self.k_proj.weight,
+                    weight2,
+                    weight2_t,
+                    weight2_s,
+                    self.v_proj.weight,
+                    weight3,
+                    weight3_t,
+                    weight3_s,
+                    rmsnorm_weight,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+        else:
+            return re_x, self.att_proj(self.attn_norm(re_x))
+
+
+class _CoatLlamaBeforeAttentionResidual(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        re_x,
+        in_x,
+        in_s,
+        weight1_origin,
+        weight1,
+        weight1_t,
+        weight1_s,
+        weight2_origin,
+        weight2,
+        weight2_t,
+        weight2_s,
+        weight3_origin,
+        weight3,
+        weight3_t,
+        weight3_s,
+        rmsnorm_weight,
+        group_size,
+        fwobits,
+        layer_id,
+        config,
+        qargs,
+        eps=1e-5,
+    ):
+        # for autograd
+        if fwobits["fabit"] == "E4M3":
+            # in_x = in_x.to(torch.float8_e4m3fn)
+            in_x = in_x.view(torch.float8_e4m3fn)
+        else:
+            raise ValueError("fabit should be E4M3")
+
+        # LayerNorm
+        ln_x, ln_s, ln_x_t, ln_utils = fp8_rmsnorm_forward(
+            in_x, in_s, rmsnorm_weight, group_size, eps, transpose_output_2d=True
+        )
+
+        # Linear Layer QKV Projection
+        if qargs.weight_memory_efficient:
+            assert weight1 is None  # memory efficient
+            weight1, weight1_s = fp8_division(weight1_origin, qargs.group_size, fwobits["fwbit"], weight1_s)
+            weight2, weight2_s = fp8_division(weight2_origin, qargs.group_size, fwobits["fwbit"], weight2_s)
+            weight3, weight3_s = fp8_division(weight3_origin, qargs.group_size, fwobits["fwbit"], weight3_s)
+
+        fc1_x = fp8_linear_forward(ln_x, ln_s, weight1, weight1_s, False, group_size)  # query states
+        fc2_x = fp8_linear_forward(ln_x, ln_s, weight2, weight2_s, False, group_size)  # key states
+        fc3_x = fp8_linear_forward(ln_x, ln_s, weight3, weight3_s, False, group_size)  # value states
+
+        # ==================== save for backward ====================
+        ctx.save_for_backward(in_x, in_s, ln_x_t, ln_s)
+        if qargs.weight_memory_efficient:
+            assert weight1_t is None and weight2_t is None and weight3_t is None
+            ctx.weight = weight1_origin, weight1_s, weight2_origin, weight2_s, weight3_origin, weight3_s
+        else:
+            ctx.weight = weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s
+
+        ctx.group_size = group_size
+        ctx.ln_utils = ln_utils
+        ctx.utils = fwobits, layer_id, config, qargs
+
+        return re_x, fc1_x, fc2_x, fc3_x
+
+    @staticmethod
+    def backward(ctx, fp_grad, query_g, key_g, value_g):
+        in_x, in_s, ln_x_t, ln_s = ctx.saved_tensors
+        weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s = ctx.weight
+
+        group_size = ctx.group_size
+        rms_weight, rstd, num_warps = ctx.ln_utils
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        # ==================== Begin backward ====================
+        # Quantize the RoPE and FlashAttention Output. grad_input and grad_weight requires different data layout.
+        query_g, query_gs, query_g_t = fp8_quantize_pertensor_transpose(
+            query_g, group_size, fwobits["babit"], transpose_output_2d=True, stochastic=False
+        )
+        key_g, key_gs, key_g_t = fp8_quantize_pertensor_transpose(
+            key_g, group_size, fwobits["babit"], transpose_output_2d=True, stochastic=False
+        )
+        value_g, value_gs, value_g_t = fp8_quantize_pertensor_transpose(
+            value_g, group_size, fwobits["babit"], transpose_output_2d=True, stochastic=False
+        )
+
+        # Linear Layer QKV Projection
+        if qargs.weight_memory_efficient:
+            weight1_t, weight1_s = fp8_division_transpose(
+                weight1_t, qargs.group_size, fwobits["fwbit"], weight1_s, only_transposed=True
+            )
+            weight2_t, weight2_s = fp8_division_transpose(
+                weight2_t, qargs.group_size, fwobits["fwbit"], weight2_s, only_transposed=True
+            )
+            weight3_t, weight3_s = fp8_division_transpose(
+                weight3_t, qargs.group_size, fwobits["fwbit"], weight3_s, only_transposed=True
+            )
+
+        fc1_g1, att_q_wg = fp8_linear_backward(
+            ln_x_t, ln_s, query_g, query_gs, query_g_t, weight1_t, weight1_s, group_size
+        )
+        fc1_g2, att_k_wg = fp8_linear_backward(ln_x_t, ln_s, key_g, key_gs, key_g_t, weight2_t, weight2_s, group_size)
+        fc1_g3, att_v_wg = fp8_linear_backward(
+            ln_x_t, ln_s, value_g, value_gs, value_g_t, weight3_t, weight3_s, group_size
+        )
+
+        fc1_g = fc1_g1 + fc1_g2 + fc1_g3
+
+        # LayerNorm
+        in_g, rms_weight_grad = fp8_rmsnorm_backward(in_x, in_s, fc1_g, rms_weight, rstd, group_size, num_warps)
+
+        # Add the gradient together, and prepare the input of the next layer.
+        re_g, (in_g, in_sg, in_sg_g16) = fp8_add_Ifp_Ifp_Ofp_Opt(
+            fp_grad, in_g, group_size, fwobits["babit"], stochastic=False
+        )
+
+        # for autograd. forward's data type should be the same of backward tensor. this will not change the actual binary representation.
+        in_g = in_g.view(torch.float8_e4m3fn)
+
+        # Although the next operator is a linear layer in MLPResidual module, we return in_sg_g16 to make the size compatible with the forward. Otherwise it will not pass autograd.
+        return (
+            re_g,
+            in_g,
+            in_sg_g16,
+            att_q_wg,
+            None,
+            None,
+            None,
+            att_k_wg,
+            None,
+            None,
+            None,
+            att_v_wg,
+            None,
+            None,
+            None,
+            rms_weight_grad,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class CoatLlamaAfterAttentionResidual(FP8CacheWeightModule):
+    """
+    This is a typical transformer attention module that contains (1) Residual (2) 1 * Linear layers
+    """
+
+    def __init__(self, config: CoatLlamaConfig, qargs: QuantizationConfig, layer_id):
+        super().__init__(config, qargs, layer_id)
+
+        self.qargs = qargs
+        self.fwobits = {
+            "fabit": self.qargs.fabit,
+            "fwbit": self.qargs.fwbit,
+            "fobit": self.qargs.fobit,
+            "babit": self.qargs.babit,
+            "bwbit": self.qargs.bwbit,
+            "bobit": self.qargs.bobit,
+        }
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
+
+    def forward(self, re_x, in_x):
+        if self.training:
+            if self.qargs.weight_memory_efficient:
+                # prepare for the weight
+                with torch.no_grad():
+                    weight4_s = self.prepare_weight(self.o_proj.weight, "o_proj", FP8Manager.is_first_microbatch)
+
+                return _CoatLlamaAfterAttentionResidual.apply(
+                    re_x,
+                    in_x,
+                    self.o_proj.weight,
+                    None,
+                    None,
+                    weight4_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+            else:
+                # prepare for the weight
+                with torch.no_grad():
+                    weight4, weight4_t, weight4_s = self.prepare_weight(
+                        self.o_proj.weight, "o_proj", FP8Manager.is_first_microbatch
+                    )
+
+                return _CoatLlamaAfterAttentionResidual.apply(
+                    re_x,
+                    in_x,
+                    self.o_proj.weight,
+                    weight4,
+                    weight4_t,
+                    weight4_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+        else:
+            return re_x + self.attn_out(in_x), None, None
+
+
+class _CoatLlamaAfterAttentionResidual(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx, re_x, flash_x, weight4_origin, weight4, weight4_t, weight4_s, group_size, fwobits, layer_id, config, qargs
+    ):
+        # Quantize the FlashAttention Output
+        flash_qx, flash_s, _ = fp8_quantize_pertensor(
+            flash_x, group_size, fwobits["fabit"]
+        )  # Modified to make it memory efficient
+
+        # # Attention Projection Linear Layer
+        if qargs.weight_memory_efficient:
+            assert weight4 is None  # memory efficient
+            weight4, weight4_s = fp8_division(weight4_origin, qargs.group_size, fwobits["fwbit"], weight4_s)
+        fc4_x = fp8_linear_forward(flash_qx, flash_s, weight4, weight4_s, False, group_size)  #
+
+        # import IPython
+        # IPython.embed()
+        # Add the activations together
+        fp_x, (out_x, out_s) = fp8_add_Ifp_Ifp_Ofp_Og16(re_x, fc4_x, flash_qx.dtype, group_size)
+
+        # ==================== save for backward ====================
+        ctx.save_for_backward(flash_x, flash_s)
+        if qargs.weight_memory_efficient:
+            assert weight4_t is None
+            ctx.weight = weight4_origin, weight4_s
+        else:
+            ctx.weight = weight4_t, weight4_s
+        ctx.group_size = group_size
+        ctx.fwobits = fwobits
+        ctx.utils = fwobits, layer_id, config, qargs
+
+        # For autograd
+        out_x = out_x.view(torch.float8_e4m3fn)
+
+        return fp_x, out_x, out_s
+
+    @staticmethod
+    def backward(ctx, fp_grad, out_g, out_gs):
+        flash_x, flash_s = ctx.saved_tensors
+        weight4_t, weight4_s = ctx.weight
+        group_size = ctx.group_size
+        fwobits = ctx.fwobits
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        # for autograd
+        if fwobits["babit"] == "E5M2":
+            # out_g = out_g.to(torch.float8_e5m2)
+            out_g = out_g.view(torch.float8_e5m2)
+        else:
+            raise ValueError("babit should be E5M2")
+        out_gs_max = out_gs.max()
+
+        # ==================== Begin backward ====================
+        # Output Projection
+        out_g_t = fp8_transpose(out_g, transpose_output_2d=True)
+
+        # We do not save an extra flash_x to save the memory usage
+        flash_x_t, flash_s = fp8_division_transpose(
+            flash_x, group_size, fwobits["fabit"], flash_s, stochastic=False, only_transposed=True
+        )
+
+        if qargs.weight_memory_efficient:
+            weight4_t, weight4_s = fp8_division_transpose(
+                weight4_t, qargs.group_size, fwobits["fwbit"], weight4_s, only_transposed=True
+            )
+        fc4_g, attn_out_wg = fp8_linear_backward(
+            flash_x_t, flash_s, out_g, out_gs_max, out_g_t, weight4_t, weight4_s, group_size
+        )
+
+        return fp_grad, fc4_g, attn_out_wg, None, None, None, None, None, None, None, None
+
+
+class CoatLlamaMLPResidual(FP8CacheWeightModule):
+    """
+    This is a typical transformer attention module that contains (1) Residual (2) LayerNorm / RMSNorm (3) 2 / 3 * Linear layers
+    (4) GELU / Silu Activation
+    """
+
+    def __init__(self, config: CoatLlamaConfig, qargs: QuantizationConfig, layer_id, hidden_size: int):
+        super().__init__(config, qargs, layer_id)
+
+        self.qargs = qargs
+        self.fwobits = {
+            "fabit": self.qargs.fabit,
+            "fwbit": self.qargs.fwbit,
+            "fobit": self.qargs.fobit,
+            "babit": self.qargs.babit,
+            "bwbit": self.qargs.bwbit,
+            "bobit": self.qargs.bobit,
+        }
+
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+        self.training = True
+
+        # below is only used when training = False
+        assert config.hidden_act == "silu", "We only support silu activation currently"
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, re_x, x, s, rmsnorm_weight):
+        if self.training:
+            if self.qargs.weight_memory_efficient:  # prepare for the weight
+                with torch.no_grad():
+                    weight1_s = self.prepare_weight(self.gate_proj.weight, "gate_proj", FP8Manager.is_first_microbatch)
+                    weight2_s = self.prepare_weight(self.up_proj.weight, "up_proj", FP8Manager.is_first_microbatch)
+                    weight3_s = self.prepare_weight(self.down_proj.weight, "down_proj", FP8Manager.is_first_microbatch)
+
+                return _CoatLlamaMLPResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.gate_proj.weight,
+                    None,
+                    None,
+                    weight1_s,
+                    self.up_proj.weight,
+                    None,
+                    None,
+                    weight2_s,
+                    self.down_proj.weight,
+                    None,
+                    None,
+                    weight3_s,
+                    rmsnorm_weight,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+            else:
+                # prepare for the weight
+                with torch.no_grad():
+                    weight1, weight1_t, weight1_s = self.prepare_weight(
+                        self.gate_proj.weight, "gate_proj", FP8Manager.is_first_microbatch
+                    )
+                    weight2, weight2_t, weight2_s = self.prepare_weight(
+                        self.up_proj.weight, "up_proj", FP8Manager.is_first_microbatch
+                    )
+                    weight3, weight3_t, weight3_s = self.prepare_weight(
+                        self.down_proj.weight, "down_proj", FP8Manager.is_first_microbatch
+                    )
+
+                return _CoatLlamaMLPResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.gate_proj.weight,
+                    weight1,
+                    weight1_t,
+                    weight1_s,
+                    self.up_proj.weight,
+                    weight2,
+                    weight2_t,
+                    weight2_s,
+                    self.down_proj.weight,
+                    weight3,
+                    weight3_t,
+                    weight3_s,
+                    rmsnorm_weight,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+        else:
+            raise NotImplementedError("Need TODO")
+            og_x = re_x
+            re_x = self.ff_norm(re_x)
+            re_x = self.ff_proj(re_x)
+            re_x = self.act(re_x)
+            re_x = self.ff_out(re_x)
+            re_x = og_x + re_x
+            return re_x, None, None
+
+
+class _CoatLlamaMLPResidual(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        re_x,
+        in_x,
+        in_s,
+        weight1_origin,
+        weight1,
+        weight1_t,
+        weight1_s,
+        weight2_origin,
+        weight2,
+        weight2_t,
+        weight2_s,
+        weight3_origin,
+        weight3,
+        weight3_t,
+        weight3_s,
+        rmsnorm_weight,
+        group_size,
+        fwobits,
+        layer_id,
+        config,
+        qargs,
+        eps=1e-5,
+    ):
+        # For autograd
+        if fwobits["fabit"] == "E4M3":
+            # in_x = in_x.to(torch.float8_e4m3fn)
+            in_x = in_x.view(torch.float8_e4m3fn)
+        else:
+            raise ValueError("fabit should be E4M3")
+
+        # LayerNorm
+        ln_x, ln_s, ln_x_t, ln_utils = fp8_rmsnorm_forward(
+            in_x, in_s, rmsnorm_weight, group_size, eps, transpose_output_2d=True
+        )
+
+        # Linear Layer of Up Projection and Gate Projection. They are fused as one linear layer.
+        if qargs.weight_memory_efficient:
+            assert weight1 is None and weight2 is None and weight3 is None  # memory efficient
+            weight1, weight1_s = fp8_division(weight1_origin, qargs.group_size, fwobits["fwbit"], weight1_s)
+            weight2, weight2_s = fp8_division(weight2_origin, qargs.group_size, fwobits["fwbit"], weight2_s)
+            weight3, weight3_s = fp8_division(weight3_origin, qargs.group_size, fwobits["fwbit"], weight3_s)
+
+        gate_x, gate_s = fp8_linear_forward(ln_x, ln_s, weight1, weight1_s, True, group_size)  # Gate Proj
+        up_x, up_s = fp8_linear_forward(ln_x, ln_s, weight2, weight2_s, True, group_size)  # Up Proj
+
+        # silu Activation
+        silu_x, silu_s = fp8_silu_forward(gate_x, gate_s, group_size)
+
+        # Element-wise Multiplication
+        mul_x, mul_s, mul_x_t = fp8_mul_forward(silu_x, silu_s, up_x, up_s, group_size, transpose_output_2d=True)
+
+        # Output Projection
+        if weight3 is None:  # memory efficient
+            weight3, weight3_s = fp8_division(weight3_origin, qargs.group_size, fwobits["fwbit"], weight3_s)
+        fc3_x = fp8_linear_forward(mul_x, mul_s, weight3, weight3_s, False, group_size)
+
+        # Add the activation together
+        fp_x, (out_x, out_s) = fp8_add_Ifp_Ifp_Ofp_Og16(re_x, fc3_x, mul_x.dtype, group_size)
+
+        # ==================== save for backward ====================
+        ctx.save_for_backward(in_x, in_s, ln_x_t, ln_s, gate_x, gate_s, up_x, up_s, silu_x, silu_s, mul_x_t, mul_s)
+
+        ctx.weight = (weight1_t, weight1_s, weight2_t, weight2_s)
+        if (
+            qargs.weight_memory_efficient
+        ):  # Weight_1/2_origin will not be saved twice, so it will be more memory efficient.
+            assert weight1_t is None and weight2_t is None and weight3_t is None
+            ctx.weight = (weight1_origin, weight1_s, weight2_origin, weight2_s, weight3_origin, weight3_s)
+        else:  # Weight1/2_t is different from the origin weight, so saving it will consumes additional memory footprint.
+            ctx.weight = (weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s)
+
+        ctx.group_size = group_size
+        ctx.ln_utils = ln_utils
+        ctx.utils = fwobits, layer_id, config, qargs
+
+        out_x = out_x.view(torch.float8_e4m3fn)
+
+        return fp_x, out_x, out_s
+
+    @staticmethod
+    def backward(ctx, fp_grad, out_g, out_gs):
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        in_x, in_s, ln_x_t, ln_s, gate_x, gate_s, up_x, up_s, silu_x, silu_s, mul_x_t, mul_s = ctx.saved_tensors
+
+        (weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s) = ctx.weight
+        group_size = ctx.group_size
+        rms_weight, rstd, num_warps = ctx.ln_utils
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        # For autograd
+        if fwobits["babit"] == "E5M2":
+            # out_g = out_g.to(torch.float8_e5m2)
+            out_g = out_g.view(torch.float8_e5m2)
+        else:
+            raise ValueError("babit should be E5M2")
+        out_gs_max = out_gs.max()
+
+        # ==================== Begin backward ====================
+        # Output Projection
+        out_gs = out_gs.max()
+        out_g_t = fp8_transpose(out_g, transpose_output_2d=True)
+
+        if qargs.weight_memory_efficient:
+            weight3_t, weight3_s = fp8_division_transpose(
+                weight3_t, qargs.group_size, fwobits["fwbit"], weight3_s, only_transposed=True
+            )
+        fc3_g, weight3_grad = fp8_linear_backward(
+            mul_x_t, mul_s, out_g, out_gs_max, out_g_t, weight3_t, weight3_s, group_size
+        )
+
+        # [MEM TEST]
+        del out_g, out_g_t, weight3_t
+
+        # Element-wise Multiplication, 1 means gate, 2 means up
+        mul_g1, (mul_g2, mul_gs2, mul_g2_t) = fp8_mul_backward(
+            silu_x, silu_s, up_x, up_s, fc3_g, group_size, fwobits["babit"], output_quantized_transpose=True
+        )
+
+        # Silu activation
+        silu_g, silu_gs, silu_g_t = fp8_silu_backward(
+            gate_x, gate_s, mul_g1, group_size, fwobits["babit"], output_quantized_transpose=True
+        )
+
+        # Linear Layer of Up and Gate Projection
+        if qargs.weight_memory_efficient:
+            weight1_t, weight1_s = fp8_division_transpose(
+                weight1_t, group_size, fwobits["fwbit"], weight1_s, only_transposed=True
+            )
+            weight2_t, weight2_s = fp8_division_transpose(
+                weight2_t, group_size, fwobits["fwbit"], weight2_s, only_transposed=True
+            )
+
+        # Gate Proj
+        fc1_g, weight1_grad = fp8_linear_backward(
+            ln_x_t, ln_s, silu_g, silu_gs, silu_g_t, weight1_t, weight1_s, group_size
+        )
+        fc2_g, weight2_grad = fp8_linear_backward(
+            ln_x_t, ln_s, mul_g2, mul_gs2, mul_g2_t, weight2_t, weight2_s, group_size
+        )
+
+        fc_g = fc1_g + fc2_g
+
+        # layerNorm
+        in_g, rms_weight_grad = fp8_rmsnorm_backward(in_x, in_s, fc_g, rms_weight, rstd, group_size, num_warps)
+
+        # Add the gradient together
+        re_g, (in_g, in_sg, in_sg_g16) = fp8_add_Ifp_Ifp_Ofp_Opt(
+            fp_grad, in_g, group_size, fwobits["babit"], stochastic=False
+        )
+
+        in_g = in_g.view(torch.float8_e4m3fn)
+
+        return (
+            re_g,
+            in_g,
+            in_sg_g16,
+            weight1_grad,
+            None,
+            None,
+            None,
+            weight2_grad,
+            None,
+            None,
+            None,
+            weight3_grad,
+            None,
+            None,
+            None,
+            rms_weight_grad,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class LlamaAttentionWithoutLinear(nn.Module):
+    """
+    Remove the Q/K/V/O projection layer in LlamaAttention module and only calculate the attention logic.
+    The Q/K/V Projection is moved to BeforeAttention Module, and the O Projection is moved to AfterAttention Module.
+    """
+
+    def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+
+        # TODO (joao): remove in v4.46 (RoPE is computed in the model, not in the decoder layers)
+        self.rotary_emb = LlamaRotaryEmbedding(config=self.config)
+
+    def forward(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = query_states.size()
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if position_embeddings is None:
+            logger.warning_once(
+                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
+                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
+                "removed and `position_embeddings` will be mandatory."
+            )
+            cos, sin = self.rotary_emb(value_states, position_ids)
+        else:
+            cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attention_mask is not None:  # no matter the length, we just slice it
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+            attn_weights = attn_weights + causal_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LlamaFlashAttention2WithoutLinear(LlamaAttentionWithoutLinear):
+    """
+    Remove the Q/K/V/O projection layer in LlamaFlashAttention2 module and only calculate the attention logic.
+    The Q/K/V Projection is moved to BeforeAttention Module, and the O Projection is moved to AfterAttention Module.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if isinstance(past_key_value, StaticCache):
+            raise ValueError(
+                "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "
+                "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers"
+            )
+
+        output_attentions = False
+
+        bsz, q_len, _ = query_states.size()
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if position_embeddings is None:
+            logger.warning_once(
+                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
+                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
+                "removed and `position_embeddings` will be mandatory."
+            )
+            cos, sin = self.rotary_emb(value_states, position_ids)
+        else:
+            cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = _flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            position_ids=position_ids,
+            dropout=dropout_rate,
+            sliding_window=getattr(self, "sliding_window", None),
+            use_top_left_mask=self._flash_attn_uses_top_left_mask,
+            is_causal=self.is_causal,
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LlamaSdpaAttentionWithoutLinear(LlamaAttentionWithoutLinear):
+    """
+    Remove the Q/K/V/O projection layer in LlamaSdpaAttention module and only calculate the attention logic.
+    The Q/K/V Projection is moved to BeforeAttention Module, and the O Projection is moved to AfterAttention Module.
+    """
+
+    # Adapted from LlamaAttention.forward
+    def forward(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "LlamaModel is using LlamaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                query_states=query_states,
+                key_states=key_states,
+                value_states=value_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+            )
+
+        bsz, q_len, _ = query_states.size()
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if position_embeddings is None:
+            logger.warning_once(
+                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
+                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
+                "removed and `position_embeddings` will be mandatory."
+            )
+            cos, sin = self.rotary_emb(value_states, position_ids)
+        else:
+            cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        causal_mask = attention_mask
+        if attention_mask is not None:
+            causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and causal_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+        # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+        is_causal = True if causal_mask is None and q_len > 1 else False
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=causal_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(bsz, q_len, -1)
+
+        return attn_output, None, past_key_value
+
+
+COAT_LLAMA_ATTENTION_CLASSES = {
+    "eager": LlamaAttentionWithoutLinear,
+    "flash_attention_2": LlamaFlashAttention2WithoutLinear,
+    "sdpa": LlamaSdpaAttentionWithoutLinear,
+}
+
+
+class CoatLlamaDecoderLayer(nn.Module):
+    def __init__(self, config: CoatLlamaConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = COAT_LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
+
+        self.qargs = QuantizationConfig(**config.coat_fp8_args)
+        self.BeforeAttention = CoatLlamaBeforeAttentionResidual(config, self.qargs, layer_idx)
+        self.AfterAttention = CoatLlamaAfterAttentionResidual(config, self.qargs, layer_idx)
+        self.MLPResidual = CoatLlamaMLPResidual(config, self.qargs, layer_idx, self.hidden_size)
+
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        quant_hidden_states: torch.Tensor,
+        scale_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): BF16 input to the layer of shape `(batch, seq_len, embed_dim)`
+            quant_hidden_states (`torch.float8_e4m3fn`): FP8 input to the layer of shape `(batch, seq_len, embed_dim)`
+            scale_hidden_states (`torch.bfloat16`): BF16 scaling factor to the layer of shape `(batch, seq_len, embed_dim // group_size)`
+            attention_mask (`torch.FloatTensor`, *optional*):
+                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+                query_sequence_length, key_sequence_length)` if default attention is used.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+                Indices depicting the position of the input sequence tokens in the sequence
+            position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+                with `head_dim` being the embedding dimension of each attention head.
+            kwargs (`dict`, *optional*):
+                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+                into the model
+        """
+
+        # Coat: The residual, LayerNorm, and the Q/K/V Projection Linear Layer
+        residual, query_states, key_states, value_states = self.BeforeAttention(
+            hidden_states, quant_hidden_states, scale_hidden_states, self.input_layernorm.weight
+        )
+
+        # Self Attention without any linear layer
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            query_states=query_states,
+            key_states=key_states,
+            value_states=value_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+
+        # Coat: The Output Projection Linear Layer and Residual
+        hidden_states, quant_hidden_states, scale_hidden_states = self.AfterAttention(residual, hidden_states)
+
+        # Residual Connection, LayerNorm, and the whole MLP module
+        hidden_states, quant_hidden_states, scale_hidden_states = self.MLPResidual(
+            hidden_states, quant_hidden_states, scale_hidden_states, self.post_attention_layernorm.weight
+        )
+
+        outputs = ((hidden_states, quant_hidden_states, scale_hidden_states),)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+class CoatLlamaPreTrainedModel(PreTrainedModel):
+    config_class = CoatLlamaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LlamaDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class CoatLlamaModel(CoatLlamaPreTrainedModel):
+    """
+    Coat Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`CoatLlamaDecoderLayer`]
+
+    Args:
+        config: CoatLlamaConfig
+    """
+
+    def __init__(self, config: CoatLlamaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [CoatLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LlamaRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Quantize
+        self.qargs = QuantizationConfig(**config.coat_fp8_args)
+        self.quantize_input_before_block = Coat_quantize_bgn(self.qargs)
+        self.quantize_output_after_block = Coat_quantize_end(self.qargs)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # kept for BC (non `Cache` `past_key_values` inputs)
+        return_legacy_cache = False
+        if use_cache and not isinstance(past_key_values, Cache):
+            return_legacy_cache = True
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+            else:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+                logger.warning_once(
+                    "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
+                    "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
+                    "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
+                )
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        # Prepare the input for Coat decoderlayer
+        hidden_states, quant_hidden_states, scale_hidden_states = self.quantize_input_before_block(hidden_states)
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    quant_hidden_states,
+                    scale_hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    quant_hidden_states,
+                    scale_hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    position_embeddings=position_embeddings,
+                )
+
+            hidden_states, quant_hidden_states, scale_hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        # Summarize the output of the Decoder Layer
+        hidden_states = self.quantize_output_after_block(hidden_states, quant_hidden_states, scale_hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if return_legacy_cache:
+            next_cache = next_cache.to_legacy_cache()
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    _update_causal_mask = LlamaModel._update_causal_mask
+
+
+class CoatLlamaForCausalLM(CoatLlamaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = CoatLlamaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    forward = LlamaForCausalLM.forward
+
+    prepare_inputs_for_generation = LlamaForCausalLM.prepare_inputs_for_generation
+
+
+# TODO
+# class LlamaForSequenceClassification(LlamaPreTrainedModel):
+
+# class LlamaForQuestionAnswering(LlamaPreTrainedModel):
+
+# class LlamaForTokenClassification(LlamaPreTrainedModel):
+
+
+def make_state_dict_compatible(state_dict: dict[str, torch.Tensor]):
+    compatible_state_dict = {}
+
+    for key, value in state_dict.items():
+        if fnmatch(key, "*self_attn.q_proj*"):
+            new_key = key.replace("self_attn.q_proj", "BeforeAttention.q_proj")
+        elif fnmatch(key, "*self_attn.k_proj*"):
+            new_key = key.replace("self_attn.k_proj", "BeforeAttention.k_proj")
+        elif fnmatch(key, "*self_attn.v_proj*"):
+            new_key = key.replace("self_attn.v_proj", "BeforeAttention.v_proj")
+        elif fnmatch(key, "*self_attn.o_proj*"):
+            new_key = key.replace("self_attn.o_proj", "AfterAttention.o_proj")
+
+        elif fnmatch(key, "*mlp.gate_proj*"):
+            new_key = key.replace("mlp.gate_proj", "MLPResidual.gate_proj")
+        elif fnmatch(key, "*mlp.up_proj*"):
+            new_key = key.replace("mlp.up_proj", "MLPResidual.up_proj")
+        elif fnmatch(key, "*mlp.down_proj*"):
+            new_key = key.replace("mlp.down_proj", "MLPResidual.down_proj")
+
+        else:
+            new_key = key
+
+        compatible_state_dict[new_key] = value
+
+    return compatible_state_dict
+
+
+AutoConfig.register("fp8_llama", CoatLlamaConfig)
+AutoModel.register(CoatLlamaConfig, CoatLlamaModel)
+AutoModelForCausalLM.register(CoatLlamaConfig, CoatLlamaForCausalLM)

llava/model/coat/activation/models/coat_llama_convert_from_hf.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+import argparse
+import os
+from dataclasses import asdict, dataclass, field
+from typing import Optional
+
+import torch
+import transformers
+from coat.activation.models._fp8_quantization_config import QuantizationConfig
+from coat.activation.models.coat_llama import CoatLlamaConfig, CoatLlamaForCausalLM, make_state_dict_compatible
+from transformers import AutoConfig, AutoModelForCausalLM
+
+
+@dataclass
+class ConvertArguments:
+    model_name: str = field(metadata={"help": "The model name or path to download the LLaMA model"})
+    save_path: str = field(metadata={"help": "The path where the converted model weights will be saved"})
+    cache_dir: str = field(default=None, metadata={"help": "Directory to cache the model"})
+
+
+def download_and_convert_llama(convert_args: ConvertArguments, quantization_args: QuantizationConfig):
+    """
+    Downloads a LLaMA model, converts its weights using `make_state_dict_compatible`,
+    and saves the converted model.
+
+    Args:
+        model_name (str): The model name or path to download the LLaMA model.
+        save_path (str): The path where the converted model weights will be saved.
+        cache_dir (Optional[str]): Directory to cache the model. Defaults to None.
+
+    Returns:
+        None
+    """
+    model_name = convert_args.model_name
+    save_path = convert_args.save_path
+    cache_dir = convert_args.cache_dir
+
+    # Step 1: Download the original LLaMA model
+    model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)
+
+    # Step 2: Initialize the model configuration for FP8 or other custom config
+    config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
+
+    # Step 3: Apply make_state_dict_compatible to convert weights
+    compatible_state_dict = make_state_dict_compatible(model.state_dict())
+
+    # Step 4: Create a new model instance with compatible configuration
+    fp8_config = CoatLlamaConfig(**config.to_dict())
+    fp8_config.coat_fp8_args = asdict(quantization_args)
+
+    converted_model = AutoModelForCausalLM.from_config(fp8_config)
+    converted_model.load_state_dict(compatible_state_dict)
+
+    # Step 5: Save the converted model and configuration using save_pretrained
+    os.makedirs(save_path, exist_ok=True)
+    converted_model.save_pretrained(save_path)
+    print(f"Converted model saved at {save_path}")
+
+
+if __name__ == "__main__":
+    # Parse command-line arguments
+    parser = transformers.HfArgumentParser((ConvertArguments, QuantizationConfig))  # NOTE: FP8
+    convert_args, quantization_args = parser.parse_args_into_dataclasses()
+
+    # Call the function with parsed arguments
+    download_and_convert_llama(convert_args, quantization_args)

llava/model/coat/activation/models/coat_olmo.py

@@ -0,0 +1,1942 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Adapted from
+[MosaiclML](https://github.com/mosaicml/examples.git) and
+[minGPT](https://github.com/karpathy/minGPT.git)
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import sys
+from abc import abstractmethod
+from collections import defaultdict
+from functools import partial
+from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Sequence, Set, Tuple, cast
+
+import torch
+import torch.backends.cuda
+import torch.nn as nn
+import torch.nn.functional as F
+from olmo.aliases import PathOrStr
+from olmo.beam_search import BeamSearch, Constraint, FinalSequenceScorer, Sampler
+from olmo.config import (
+    ActivationCheckpointingStrategy,
+    ActivationType,
+    BlockType,
+    CheckpointType,
+    FSDPWrapStrategy,
+    InitFnType,
+    LayerNormType,
+    ModelConfig,
+    QuantActivationConfig,
+    ShardedCheckpointerType,
+    TrainConfig,
+)
+from olmo.exceptions import OLMoConfigurationError
+from olmo.initialization import init_normal
+from olmo.model import (
+    Activation,
+    BufferCache,
+    Dropout,
+    LayerNorm,
+    LayerNormBase,
+    OLMo,
+    OLMoBlock,
+    OLMoBlockGroup,
+    OLMoGenerateOutput,
+    OLMoOutput,
+    RMSLayerNorm,
+    RotaryEmbedding,
+    _non_meta_init_device,
+    activation_checkpoint_function,
+    alibi_attention_bias,
+    causal_attention_bias,
+    get_causal_attention_bias,
+    should_checkpoint_block,
+)
+from olmo.torch_util import ensure_finite_, get_cumulative_document_lengths
+from torch import einsum
+
+from ..real_quantization import (
+    Coat_quantize_bgn,
+    Coat_quantize_end,
+    fp8_add_Ifp_Ifp_Ofp_Og16,
+    fp8_add_Ifp_Ifp_Ofp_Opt,
+    fp8_division,
+    fp8_division_transpose,
+    fp8_gelu_backward,
+    fp8_gelu_forward,
+    fp8_layernorm_noparam_backward,
+    fp8_layernorm_noparam_forward,
+    fp8_linear_backward,
+    fp8_linear_forward,
+    fp8_mul_backward,
+    fp8_mul_forward,
+    fp8_quantize,
+    fp8_quantize_pertensor,
+    fp8_quantize_pertensor_transpose,
+    fp8_rmsnorm_backward,
+    fp8_rmsnorm_forward,
+    fp8_silu_backward,
+    fp8_silu_forward,
+    fp8_transpose,
+)
+from ._fp8_weightcache import FP8CacheWeightModule
+from ._fp8manager import FP8Manager
+
+if sys.version_info.minor > 8:
+    from collections.abc import MutableMapping
+elif sys.version_info.minor == 8:
+    from typing import MutableMapping
+else:
+    raise SystemExit("This script supports Python 3.8 or higher")
+
+__all__ = [
+    "LayerNormBase",
+    "LayerNorm",
+    "RMSLayerNorm",
+    "RotaryEmbedding",
+    "Activation",
+    "GELU",
+    "ReLU",
+    "SwiGLU",
+    "OLMoBlock",
+    "OLMoSequentialBlock",
+    "OLMo",
+    "OLMoOutput",
+    "OLMoGenerateOutput",
+]
+
+
+log = logging.getLogger(__name__)
+
+
+class CoatOLMoBeforeAttentionResidual(FP8CacheWeightModule):
+    """
+    This is a typical transformer attention module that contains (1) Residual (2) LayerNorm / RMSNorm (3) 1 * Linear layers
+    """
+
+    def __init__(self, config: ModelConfig, qargs: QuantActivationConfig, layer_id, fused_dims: tuple):
+        super().__init__(config, qargs, layer_id)
+
+        self.qargs = qargs
+        self.fwobits = {
+            "fabit": self.qargs.fabit,
+            "fwbit": self.qargs.fwbit,
+            "fobit": self.qargs.fobit,
+            "babit": self.qargs.babit,
+            "bwbit": self.qargs.bwbit,
+            "bobit": self.qargs.bobit,
+        }
+        self.ln_normalized_shape = config.d_model
+        self.att_proj = nn.Linear(config.d_model, sum(fused_dims), bias=config.include_bias, device=config.init_device)
+
+        self.attn_norm = LayerNorm.build(config)
+
+    def forward(self, re_x, x, s):
+        if self.training:
+            if self.qargs.weight_memory_efficient:
+                # Prepare
+                with torch.no_grad():
+                    weight1_s = self.prepare_weight(self.att_proj.weight, "att_proj", FP8Manager.is_first_microbatch)
+                return _CoatOLMoBeforeAttentionResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.att_proj.weight,
+                    None,
+                    None,
+                    weight1_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+            else:
+                # Prepare
+                with torch.no_grad():
+                    weight1, weight1_t, weight1_s = self.prepare_weight(
+                        self.att_proj.weight, "att_proj", FP8Manager.is_first_microbatch
+                    )
+                return _CoatOLMoBeforeAttentionResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.att_proj.weight,
+                    weight1,
+                    weight1_t,
+                    weight1_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+        else:
+            return re_x, self.att_proj(self.attn_norm(re_x))
+
+
+class _CoatOLMoBeforeAttentionResidual(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        re_x,
+        in_x,
+        in_s,
+        weight1_origin,
+        weight1,
+        weight1_t,
+        weight1_s,
+        group_size,
+        fwobits,
+        layer_id,
+        config,
+        qargs,
+        eps=1e-5,
+    ):
+        # for autograd
+        if fwobits["fabit"] == "E4M3":
+            # in_x = in_x.to(torch.float8_e4m3fn)
+            in_x = in_x.view(torch.float8_e4m3fn)
+        else:
+            raise ValueError("fabit should be E4M3")
+
+        # LayerNorm
+        ln_x, ln_s, ln_x_t, ln_utils = fp8_layernorm_noparam_forward(
+            in_x, in_s, group_size, eps, transpose_output_2d=True
+        )
+
+        # Linear Layer QKV Projection
+        if qargs.weight_memory_efficient:
+            assert weight1 is None  # memory efficient
+            weight1, weight1_s = fp8_division(weight1_origin, qargs.group_size, fwobits["fwbit"], weight1_s)
+        fc1_x = fp8_linear_forward(ln_x, ln_s, weight1, weight1_s, False, group_size)
+
+        # ==================== save for backward ====================
+        ctx.save_for_backward(in_x, in_s, ln_x_t, ln_s)
+        if qargs.weight_memory_efficient:
+            assert weight1_t is None
+            ctx.weight = weight1_origin, weight1_s
+        else:
+            ctx.weight = weight1_t, weight1_s
+        ctx.group_size = group_size
+        ctx.ln_utils = ln_utils
+        ctx.utils = fwobits, layer_id, config, qargs
+
+        return re_x, fc1_x
+
+    @staticmethod
+    def backward(ctx, fp_grad, flash_g):
+        in_x, in_s, ln_x_t, ln_s = ctx.saved_tensors
+        weight1_t, weight1_s = ctx.weight
+        group_size = ctx.group_size
+        mean, rstd, num_warps = ctx.ln_utils
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        # ==================== Begin backward ====================
+        # Quantize the RoPE and FlashAttention Output. grad_input and grad_weight requires different data layout.
+        flash_g, flash_gs, flash_g_t = fp8_quantize_pertensor_transpose(
+            flash_g, group_size, fwobits["babit"], transpose_output_2d=True, stochastic=False
+        )
+
+        # Linear Layer QKV Projection
+        if qargs.weight_memory_efficient:
+            weight1_t, weight1_s = fp8_division_transpose(
+                weight1_t, qargs.group_size, fwobits["fwbit"], weight1_s, only_transposed=True
+            )
+        fc1_g, att_proj_wg = fp8_linear_backward(
+            ln_x_t, ln_s, flash_g, flash_gs, flash_g_t, weight1_t, weight1_s, group_size
+        )
+
+        # LayerNorm
+        in_g = fp8_layernorm_noparam_backward(in_x, in_s, fc1_g, group_size, mean, rstd, num_warps)
+
+        # Add the gradient together, and prepare the input of the next layer.
+        re_g, (in_g, in_sg, in_sg_g16) = fp8_add_Ifp_Ifp_Ofp_Opt(
+            fp_grad, in_g, group_size, fwobits["babit"], stochastic=False
+        )
+
+        # for autograd. forward's data type should be the same of backward tensor. this will not change the actual binary representation.
+        in_g = in_g.view(torch.float8_e4m3fn)
+
+        # Although the next operator is a linear layer in MLPResidual module, we return in_sg_g16 to make the size compatible with the forward. Otherwise it will not pass autograd.
+        return re_g, in_g, in_sg_g16, att_proj_wg, None, None, None, None, None, None, None, None, None
+
+
+class CoatOLMoAfterAttentionResidual(FP8CacheWeightModule):
+    """
+    This is a typical transformer attention module that contains (1) Residual (2) 1 * Linear layers
+    """
+
+    def __init__(self, config: ModelConfig, qargs: QuantActivationConfig, layer_id):
+        super().__init__(config, qargs, layer_id)
+
+        self.qargs = qargs
+        self.fwobits = {
+            "fabit": self.qargs.fabit,
+            "fwbit": self.qargs.fwbit,
+            "fobit": self.qargs.fobit,
+            "babit": self.qargs.babit,
+            "bwbit": self.qargs.bwbit,
+            "bobit": self.qargs.bobit,
+        }
+        self.attn_out = nn.Linear(config.d_model, config.d_model, bias=config.include_bias, device=config.init_device)
+
+    def forward(self, re_x, in_x):
+        if self.training:
+            if self.qargs.weight_memory_efficient:
+                # prepare for the weight
+                with torch.no_grad():
+                    weight2_s = self.prepare_weight(self.attn_out.weight, "attn_out", FP8Manager.is_first_microbatch)
+
+                return _CoatOLMoAfterAttentionResidual.apply(
+                    re_x,
+                    in_x,
+                    self.attn_out.weight,
+                    None,
+                    None,
+                    weight2_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+            else:
+                # prepare for the weight
+                with torch.no_grad():
+                    weight2, weight2_t, weight2_s = self.prepare_weight(
+                        self.attn_out.weight, "attn_out", FP8Manager.is_first_microbatch
+                    )
+
+                return _CoatOLMoAfterAttentionResidual.apply(
+                    re_x,
+                    in_x,
+                    self.attn_out.weight,
+                    weight2,
+                    weight2_t,
+                    weight2_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+        else:
+            return re_x + self.attn_out(in_x), None, None
+
+
+class _CoatOLMoAfterAttentionResidual(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx, re_x, flash_x, weight2_origin, weight2, weight2_t, weight2_s, group_size, fwobits, layer_id, config, qargs
+    ):
+        # Quantize the FlashAttention Output
+        flash_qx, flash_s, _ = fp8_quantize_pertensor(
+            flash_x, group_size, fwobits["fabit"]
+        )  # Modified to make it memory efficient
+
+        # # Attention Projection Linear Layer
+        if qargs.weight_memory_efficient:
+            assert weight2 is None  # memory efficient
+            weight2, weight2_s = fp8_division(weight2_origin, qargs.group_size, fwobits["fwbit"], weight2_s)
+        fc2_x = fp8_linear_forward(flash_qx, flash_s, weight2, weight2_s, False, group_size)  #
+
+        # import IPython
+        # IPython.embed()
+        # Add the activations together
+        fp_x, (out_x, out_s) = fp8_add_Ifp_Ifp_Ofp_Og16(re_x, fc2_x, flash_qx.dtype, group_size)
+
+        # ==================== save for backward ====================
+        ctx.save_for_backward(flash_x, flash_s)
+        if qargs.weight_memory_efficient:
+            assert weight2_t is None
+            ctx.weight = weight2_origin, weight2_s
+        else:
+            ctx.weight = weight2_t, weight2_s
+        ctx.group_size = group_size
+        ctx.fwobits = fwobits
+        ctx.utils = fwobits, layer_id, config, qargs
+
+        # For autograd
+        out_x = out_x.view(torch.float8_e4m3fn)
+
+        return fp_x, out_x, out_s
+
+    @staticmethod
+    def backward(ctx, fp_grad, out_g, out_gs):
+        flash_x, flash_s = ctx.saved_tensors
+        weight2_t, weight2_s = ctx.weight
+        group_size = ctx.group_size
+        fwobits = ctx.fwobits
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        # for autograd
+        if fwobits["babit"] == "E5M2":
+            # out_g = out_g.to(torch.float8_e5m2)
+            out_g = out_g.view(torch.float8_e5m2)
+        else:
+            raise ValueError("babit should be E5M2")
+        out_gs_max = out_gs.max()
+
+        # ==================== Begin backward ====================
+        # Output Projection
+        out_g_t = fp8_transpose(out_g, transpose_output_2d=True)
+
+        # We do not save an extra flash_x to save the memory usage
+        flash_x_t, flash_s = fp8_division_transpose(
+            flash_x, group_size, fwobits["fabit"], flash_s, stochastic=False, only_transposed=True
+        )
+
+        if qargs.weight_memory_efficient:
+            weight2_t, weight2_s = fp8_division_transpose(
+                weight2_t, qargs.group_size, fwobits["fwbit"], weight2_s, only_transposed=True
+            )
+        fc2_g, attn_out_wg = fp8_linear_backward(
+            flash_x_t, flash_s, out_g, out_gs_max, out_g_t, weight2_t, weight2_s, group_size
+        )
+
+        return fp_grad, fc2_g, attn_out_wg, None, None, None, None, None, None, None, None
+
+
+class CoatOLMoMLPResidual(FP8CacheWeightModule):
+    """
+    This is a typical transformer attention module that contains (1) Residual (2) LayerNorm / RMSNorm (3) 2 / 3 * Linear layers
+    (4) GELU / Silu Activation
+    """
+
+    def __init__(self, config: ModelConfig, qargs: QuantActivationConfig, layer_id, hidden_size: int):
+        super().__init__(config, qargs, layer_id)
+
+        self.qargs = qargs
+        self.fwobits = {
+            "fabit": self.qargs.fabit,
+            "fwbit": self.qargs.fwbit,
+            "fobit": self.qargs.fobit,
+            "babit": self.qargs.babit,
+            "bwbit": self.qargs.bwbit,
+            "bobit": self.qargs.bobit,
+        }
+        self.ln_normalized_shape = config.d_model
+        self.act_output_multiplier = 0.5 if config.activation_type == ActivationType.swiglu else 1
+        self.ff_proj = nn.Linear(config.d_model, hidden_size, bias=config.include_bias, device=config.init_device)
+        self.ff_out = nn.Linear(
+            int(self.act_output_multiplier * hidden_size),
+            config.d_model,
+            bias=config.include_bias,
+            device=config.init_device,
+        )
+        self.training = True
+
+        # below is only used when training = False
+        self.ff_norm = LayerNorm.build(config)
+        self.act = Activation.build(config)
+        assert (self.act.output_multiplier * hidden_size) % 1 == 0
+
+    def forward(self, re_x, x, s):
+        if self.training:
+            if self.qargs.weight_memory_efficient:  # prepare for the weight
+                with torch.no_grad():
+                    weight1_s = self.prepare_weight(self.ff_proj.weight, "ff_proj", FP8Manager.is_first_microbatch)
+                    weight2_s = self.prepare_weight(self.ff_out.weight, "ff_out", FP8Manager.is_first_microbatch)
+
+                return _CoatOLMoMLPResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.ff_proj.weight,
+                    None,
+                    None,
+                    weight1_s,
+                    self.ff_out.weight,
+                    None,
+                    None,
+                    weight2_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+            else:
+                # prepare for the weight
+                with torch.no_grad():
+                    weight1, weight1_t, weight1_s = self.prepare_weight(
+                        self.ff_proj.weight, "ff_proj", FP8Manager.is_first_microbatch
+                    )
+                    weight2, weight2_t, weight2_s = self.prepare_weight(
+                        self.ff_out.weight, "ff_out", FP8Manager.is_first_microbatch
+                    )
+
+                return _CoatOLMoMLPResidual.apply(
+                    re_x,
+                    x,
+                    s,
+                    self.ff_proj.weight,
+                    weight1,
+                    weight1_t,
+                    weight1_s,
+                    self.ff_out.weight,
+                    weight2,
+                    weight2_t,
+                    weight2_s,
+                    self.qargs.group_size,
+                    self.fwobits,
+                    self.layer_id,
+                    self.config,
+                    self.qargs,
+                )
+        else:
+            og_x = re_x
+            re_x = self.ff_norm(re_x)
+            re_x = self.ff_proj(re_x)
+            re_x = self.act(re_x)
+            re_x = self.ff_out(re_x)
+            re_x = og_x + re_x
+            return re_x, None, None
+
+
+class _CoatOLMoMLPResidual(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        re_x,
+        in_x,
+        in_s,
+        weight1_origin,
+        weight1,
+        weight1_t,
+        weight1_s,
+        weight2_origin,
+        weight2,
+        weight2_t,
+        weight2_s,
+        group_size,
+        fwobits,
+        layer_id,
+        config,
+        qargs,
+        eps=1e-5,
+    ):
+        # For autograd
+        if fwobits["fabit"] == "E4M3":
+            # in_x = in_x.to(torch.float8_e4m3fn)
+            in_x = in_x.view(torch.float8_e4m3fn)
+        else:
+            raise ValueError("fabit should be E4M3")
+
+        # LayerNorm
+        ln_x, ln_s, ln_x_t, ln_utils = fp8_layernorm_noparam_forward(
+            in_x, in_s, group_size, eps, transpose_output_2d=True
+        )
+
+        # Linear Layer of Up Projection and Gate Projection. They are fused as one linear layer.
+        if qargs.weight_memory_efficient:
+            assert weight1 is None  # memory efficient
+            weight1, weight1_s = fp8_division(weight1_origin, qargs.group_size, fwobits["fwbit"], weight1_s)
+        fc1_x, fc1_s = fp8_linear_forward(ln_x, ln_s, weight1, weight1_s, True, group_size)
+
+        # NOTE: Becareful of the order
+        up_x, gate_x = fc1_x.chunk(2, dim=-1)
+        up_s, gate_s = fc1_s.chunk(2, dim=-1)
+
+        # silu Activation
+        silu_x, silu_s = fp8_silu_forward(gate_x, gate_s, group_size)
+
+        # Element-wise Multiplication
+        mul_x, mul_s, mul_x_t = fp8_mul_forward(silu_x, silu_s, up_x, up_s, group_size, transpose_output_2d=True)
+
+        # Output Projection
+        if weight2 is None:  # memory efficient
+            weight2, weight2_s = fp8_division(weight2_origin, qargs.group_size, fwobits["fwbit"], weight2_s)
+        fc2_x = fp8_linear_forward(mul_x, mul_s, weight2, weight2_s, False, group_size)
+
+        # Add the activation together
+        fp_x, (out_x, out_s) = fp8_add_Ifp_Ifp_Ofp_Og16(re_x, fc2_x, mul_x.dtype, group_size)
+
+        # ==================== save for backward ====================
+        ctx.save_for_backward(in_x, in_s, ln_x_t, ln_s, gate_x, gate_s, up_x, up_s, silu_x, silu_s, mul_x_t, mul_s)
+
+        ctx.weight = (weight1_t, weight1_s, weight2_t, weight2_s)
+        if (
+            qargs.weight_memory_efficient
+        ):  # Weight_1/2_origin will not be saved twice, so it will be more memory efficient.
+            assert weight1_t is None
+            ctx.weight = (weight1_origin, weight1_s, weight2_origin, weight2_s)
+        else:  # Weight1/2_t is different from the origin weight, so saving it will consumes additional memory footprint.
+            ctx.weight = (weight1_t, weight1_s, weight2_t, weight2_s)
+
+        ctx.group_size = group_size
+        ctx.ln_utils = ln_utils
+        ctx.utils = fwobits, layer_id, config, qargs
+
+        out_x = out_x.view(torch.float8_e4m3fn)
+
+        return fp_x, out_x, out_s
+
+    @staticmethod
+    def backward(ctx, fp_grad, out_g, out_gs):
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        in_x, in_s, ln_x_t, ln_s, gate_x, gate_s, up_x, up_s, silu_x, silu_s, mul_x_t, mul_s = ctx.saved_tensors
+
+        (weight1_t, weight1_s, weight2_t, weight2_s) = ctx.weight
+        group_size = ctx.group_size
+        mean, rstd, num_warps = ctx.ln_utils
+        fwobits, layer_id, config, qargs = ctx.utils
+
+        # For autograd
+        if fwobits["babit"] == "E5M2":
+            # out_g = out_g.to(torch.float8_e5m2)
+            out_g = out_g.view(torch.float8_e5m2)
+        else:
+            raise ValueError("babit should be E5M2")
+        out_gs_max = out_gs.max()
+
+        # ==================== Begin backward ====================
+        # Output Projection
+        out_gs = out_gs.max()
+        out_g_t = fp8_transpose(out_g, transpose_output_2d=True)
+
+        if qargs.weight_memory_efficient:
+            weight2_t, weight2_s = fp8_division_transpose(
+                weight2_t, qargs.group_size, fwobits["fwbit"], weight2_s, only_transposed=True
+            )
+        fc2_g, weight2_grad = fp8_linear_backward(
+            mul_x_t, mul_s, out_g, out_gs_max, out_g_t, weight2_t, weight2_s, group_size
+        )
+
+        # [MEM TEST]
+        del out_g, out_g_t, weight2_t
+
+        # Element-wise Multiplication, 1 means gate, 2 means up
+        mul_g1, (mul_g2, mul_gs2) = fp8_mul_backward(silu_x, silu_s, up_x, up_s, fc2_g, group_size, fwobits["babit"])
+
+        # Silu activation
+        silu_g, silu_gs = fp8_silu_backward(gate_x, gate_s, mul_g1, group_size, fwobits["babit"])
+
+        # Prepare the input of Linear Layer. NOTE: Becareful of the order
+        gateup_g = torch.cat([mul_g2, silu_g], dim=-1)
+        gateup_gs = torch.cat([mul_gs2, silu_gs])
+        gateup_gs = torch.max(gateup_gs)
+
+        gateup_g, gateup_gs, gateup_g_t = fp8_division_transpose(
+            gateup_g, group_size, fwobits["babit"], gateup_gs, stochastic=False
+        )
+
+        # Linear Layer of Up and Gate Projection
+        if qargs.weight_memory_efficient:
+            weight1_t, weight1_s = fp8_division_transpose(
+                weight1_t, group_size, fwobits["fwbit"], weight1_s, only_transposed=True
+            )
+        fc1_g, weight1_grad = fp8_linear_backward(
+            ln_x_t, ln_s, gateup_g, gateup_gs, gateup_g_t, weight1_t, weight1_s, group_size
+        )
+
+        # layerNorm
+        in_g = fp8_layernorm_noparam_backward(in_x, in_s, fc1_g, group_size, mean, rstd, num_warps)
+
+        # Add the gradient together
+        re_g, (in_g, in_sg, in_sg_g16) = fp8_add_Ifp_Ifp_Ofp_Opt(
+            fp_grad, in_g, group_size, fwobits["babit"], stochastic=False
+        )
+
+        in_g = in_g.view(torch.float8_e4m3fn)
+
+        return (
+            re_g,
+            in_g,
+            in_sg_g16,
+            weight1_grad,
+            None,
+            None,
+            None,
+            weight2_grad,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class CoatOLMoBlock(nn.Module):
+    """
+    A base class for transformer block implementations.
+    """
+
+    def __init__(self, layer_id: int, config: ModelConfig, qargs: QuantActivationConfig, cache: BufferCache):
+        super().__init__()
+        self.layer_id = layer_id
+        self.config = config
+        self.qargs = qargs
+        self.hidden_size = (
+            config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model
+        )
+        self.__cache = cache
+        assert config.d_model % config.n_heads == 0
+
+        self._activation_checkpoint_fn: Callable | None = None
+
+        # Dropout.
+        self.dropout = Dropout(config.residual_dropout)
+
+        # Layer norms.
+        self.k_norm: LayerNormBase | None = None
+        self.q_norm: LayerNormBase | None = None
+        if config.attention_layer_norm:
+            assert config.effective_n_kv_heads is not None
+            self.k_norm = LayerNormBase.build(
+                config,
+                size=(config.d_model // config.n_heads) * config.effective_n_kv_heads,
+                elementwise_affine=config.attention_layer_norm_with_affine,
+            )
+            self.q_norm = LayerNormBase.build(config, elementwise_affine=config.attention_layer_norm_with_affine)
+
+        # Make sure QKV clip coefficient is positive, otherwise it's not well-defined.
+        if config.clip_qkv is not None:
+            assert config.clip_qkv > 0
+
+        # Activation function.
+        self.act = Activation.build(config)
+        assert (self.act.output_multiplier * self.hidden_size) % 1 == 0
+
+        if not self.qargs.use_quantize_model:
+            # Attention output projection.
+            self.attn_out = nn.Linear(
+                config.d_model, config.d_model, bias=config.include_bias, device=config.init_device
+            )
+
+            # Feed-forward output projection.
+            self.ff_out = nn.Linear(
+                int(self.act.output_multiplier * self.hidden_size),
+                config.d_model,
+                bias=config.include_bias,
+                device=config.init_device,
+            )
+            self.ff_out._is_residual = True  # type: ignore
+
+        # Rotary embeddings.
+        if self.config.rope:
+            self.rotary_emb = RotaryEmbedding(config, self.__cache)
+
+        self.flash_attn_func = None
+        self.flash_attn_varlen_func = None
+        if config.flash_attention:
+            try:
+                from flash_attn import flash_attn_func, flash_attn_varlen_func  # type: ignore
+
+                self.flash_attn_func = flash_attn_func
+                self.flash_attn_varlen_func = flash_attn_varlen_func
+            except ModuleNotFoundError:
+                pass
+
+    def reset_parameters(self):
+        if self.k_norm is not None:
+            self.k_norm.reset_parameters()
+        if self.q_norm is not None:
+            self.q_norm.reset_parameters()
+
+        if not self.qargs.use_quantize_model:
+            if self.config.init_fn == InitFnType.normal:
+                attn_out_std = ff_out_std = self.config.init_std
+                cutoff_factor = self.config.init_cutoff_factor
+
+            elif self.config.init_fn == InitFnType.mitchell:
+                attn_out_std = 1 / (math.sqrt(2 * self.config.d_model * (self.layer_id + 1)))
+                ff_out_std = 1 / (math.sqrt(2 * self.ff_out.in_features * (self.layer_id + 1)))
+                cutoff_factor = self.config.init_cutoff_factor or 3.0
+
+            elif self.config.init_fn == InitFnType.full_megatron:
+                attn_out_std = ff_out_std = self.config.init_std / math.sqrt(2.0 * self.config.n_layers)
+                cutoff_factor = self.config.init_cutoff_factor or 3.0
+
+            else:
+                raise NotImplementedError(self.config.init_fn)
+
+            init_normal(self.attn_out, std=attn_out_std, init_cutoff_factor=cutoff_factor)
+            init_normal(self.ff_out, std=ff_out_std, init_cutoff_factor=cutoff_factor)
+
+    def set_activation_checkpointing(
+        self, strategy: ActivationCheckpointingStrategy | None, checkpoint_func: Callable | None = None
+    ):
+        if strategy == ActivationCheckpointingStrategy.fine_grained:
+            self._activation_checkpoint_fn = checkpoint_func or activation_checkpoint_function(self.config)
+        else:
+            self._activation_checkpoint_fn = None
+
+    @classmethod
+    def _cast_attn_bias(cls, bias: torch.Tensor, input_dtype: torch.dtype) -> torch.Tensor:
+        target_dtype = input_dtype
+        # NOTE: `is_autocast_enabled()` only checks for CUDA autocast, so we use the separate function
+        # `is_autocast_cpu_enabled()` for CPU autocast.
+        # See https://github.com/pytorch/pytorch/issues/110966.
+        if bias.device.type == "cuda" and torch.is_autocast_enabled():
+            target_dtype = torch.get_autocast_gpu_dtype()
+        elif bias.device.type == "cpu" and torch.is_autocast_cpu_enabled():
+            target_dtype = torch.get_autocast_cpu_dtype()
+        if bias.dtype != target_dtype:
+            bias = bias.to(target_dtype)
+            ensure_finite_(bias, check_neg_inf=True, check_pos_inf=False)
+        return bias
+
+    def _scaled_dot_product_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        max_doc_len: int | None = None,
+        cu_doc_lens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Computes scaled dot product attention on query, key and value tensors, using an optional
+        attention mask if passed, and applying dropout if a probability greater than 0.0 is specified.
+        """
+        if max_doc_len is not None and cu_doc_lens is not None:
+            assert self.flash_attn_varlen_func is not None, "flash-attn is required for document masking"
+            assert attn_mask is None, "attn-mask is currently not supported with document masking"
+            B, T, D = q.size(0), q.size(2), q.size(3)
+            r = self.flash_attn_varlen_func(
+                q.transpose(1, 2).view(B * T, -1, D),
+                k.transpose(1, 2).view(B * T, -1, D),
+                v.transpose(1, 2).view(B * T, -1, D),
+                cu_doc_lens,
+                cu_doc_lens,
+                max_doc_len,
+                max_doc_len,
+                dropout_p=dropout_p,
+                causal=is_causal,
+            )
+            return r.view(B, T, -1, D).transpose(1, 2)
+        elif self.flash_attn_func is not None and attn_mask is None:
+            r = self.flash_attn_func(
+                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p, causal=is_causal
+            )
+            return r.transpose(1, 2)
+        else:
+            # torch's sdpa doesn't support GQA, so we're doing this
+            assert k.size(1) == v.size(1)
+            num_kv_heads = k.size(1)
+            num_q_heads = q.size(1)
+            if num_q_heads != num_kv_heads:
+                assert num_q_heads % num_kv_heads == 0
+                k = k.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads)
+                v = v.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads)
+
+            return F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attn_mask,
+                dropout_p=dropout_p,
+                is_causal=is_causal,
+            )
+
+    def attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attention_bias: torch.Tensor | None = None,
+        layer_past: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+        max_doc_len: int | None = None,
+        cu_doc_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        B, T, C = q.size()  # batch size, sequence length, d_model
+        dtype = k.dtype
+
+        # Optionally apply layer norm to keys and queries.
+        if self.q_norm is not None and self.k_norm is not None:
+            q = self.q_norm(q).to(dtype=dtype)
+            k = self.k_norm(k).to(dtype=dtype)
+
+        # Move head forward to be next to the batch dim.
+        # shape: (B, nh, T, hs)
+        q = q.view(B, T, self.config.n_heads, C // self.config.n_heads).transpose(1, 2)
+        # shape: (B, n_kv_h, T, hs)
+        k = k.view(B, T, self.config.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2)
+        # shape: (B, n_kv_h, T, hs)
+        v = v.view(B, T, self.config.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            k = torch.cat((past_key, k), dim=-2)
+            v = torch.cat((past_value, v), dim=-2)
+
+        present = (k, v) if use_cache else None
+        query_len, key_len = q.shape[-2], k.shape[-2]  # could be different if layer_past not None
+
+        if self.config.rope:
+            # Apply rotary embeddings.
+            q, k = self.rotary_emb(q, k)
+
+        if attention_bias is not None:
+            # Resize and cast attention bias.
+            # The current dtype of the attention bias might not match the dtype that the SDP attn function will
+            # run in if AMP is enabled, and this can be a problem if some tokens are masked out due to padding
+            # as down-casting the attention bias to the autocast precision will result in -infs, which will
+            # cause the SDP attn function to produce NaNs.
+            attention_bias = self._cast_attn_bias(attention_bias[:, :, key_len - query_len : key_len, :key_len], dtype)
+
+        # Get the attention scores.
+        # shape: (B, nh, T, hs)
+        att = self._scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attention_bias,
+            dropout_p=0.0 if not self.training else self.config.attention_dropout,
+            is_causal=attention_bias is None,
+            max_doc_len=max_doc_len,
+            cu_doc_lens=cu_doc_lens,
+        )
+
+        # Re-assemble all head outputs side-by-side.
+        att = att.transpose(1, 2).contiguous().view(B, T, C)
+
+        # Apply output projection. NOTE: We move the attn output outside of this attention function
+        return att, present
+
+    @abstractmethod
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_bias: torch.FloatTensor | None = None,
+        layer_past: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+        max_doc_len: int | None = None,
+        cu_doc_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        raise NotImplementedError
+
+    @classmethod
+    def build(cls, layer_id: int, config: ModelConfig, qargs: QuantActivationConfig, cache: BufferCache) -> OLMoBlock:
+        if config.block_type == BlockType.sequential:
+            return CoatOLMoSequentialBlock(layer_id, config, qargs, cache)
+        elif config.block_type == BlockType.llama:
+            return CoatOLMoLlamaBlock(layer_id, config, qargs, cache)
+        else:
+            raise NotImplementedError(f"Unknown block type: '{config.block_type}'")
+
+
+class CoatOLMoSequentialBlock(CoatOLMoBlock):
+    """
+    This is a typical transformer block where the output is computed as ``MLP(LN(x + Attention(LN(x))))``
+    (plus another skip connection). To compute it as ``LN(MLP(x + LN(Attention(x))))``,
+    use the flag `norm_after`.
+    """
+
+    def __init__(self, layer_id: int, config: ModelConfig, qargs: QuantActivationConfig, cache: BufferCache):
+        super().__init__(layer_id, config, qargs, cache)
+        # Attention input projection. Projects x -> (q, k, v)
+
+        assert not self.config.norm_after, "COAT currently does not support PostNorm"
+
+        head_dim = config.d_model // config.n_heads
+        self.fused_dims = (
+            config.d_model,
+            config.effective_n_kv_heads * head_dim,
+            config.effective_n_kv_heads * head_dim,
+        )
+
+        if self.qargs.use_quantize_model:
+            self.BeforeAttention = CoatOLMoBeforeAttentionResidual(config, qargs, self.layer_id, self.fused_dims)
+            self.AfterAttention = CoatOLMoAfterAttentionResidual(config, qargs, self.layer_id)
+            self.MLPResidual = CoatOLMoMLPResidual(config, qargs, self.layer_id, self.hidden_size)
+        else:
+            self.att_proj = nn.Linear(
+                config.d_model, sum(self.fused_dims), bias=config.include_bias, device=config.init_device
+            )
+            # Feed-forward input projection.
+            self.ff_proj = nn.Linear(
+                config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device
+            )
+
+        # Layer norms.
+        self.attn_norm = LayerNorm.build(config, size=config.d_model)
+        self.ff_norm = LayerNorm.build(config, size=config.d_model)
+
+    def reset_parameters(self):
+        super().reset_parameters()
+        self.attn_norm.reset_parameters()
+        self.ff_norm.reset_parameters()
+        # NOTE: the standard deviation for these weights does not depend on the layer.
+
+        if self.qargs.use_quantize_model:  # The initialization appears here, not in CoatOLMoBlock's reset_parameters
+            if self.config.init_fn == InitFnType.normal:
+                attn_out_std = ff_out_std = self.config.init_std
+                cutoff_factor = self.config.init_cutoff_factor
+
+            elif self.config.init_fn == InitFnType.mitchell:
+                attn_out_std = 1 / (math.sqrt(2 * self.config.d_model * (self.layer_id + 1)))
+                ff_out_std = 1 / (math.sqrt(2 * self.MLPResidual.ff_out.in_features * (self.layer_id + 1)))
+                cutoff_factor = self.config.init_cutoff_factor or 3.0
+
+            elif self.config.init_fn == InitFnType.full_megatron:
+                attn_out_std = ff_out_std = self.config.init_std / math.sqrt(2.0 * self.config.n_layers)
+                cutoff_factor = self.config.init_cutoff_factor or 3.0
+
+            else:
+                raise NotImplementedError(self.config.init_fn)
+
+            init_normal(self.AfterAttention.attn_out, std=attn_out_std, init_cutoff_factor=cutoff_factor)
+            init_normal(self.MLPResidual.ff_out, std=ff_out_std, init_cutoff_factor=cutoff_factor)
+
+        if self.config.init_fn == InitFnType.normal:
+            std = self.config.init_std
+            cutoff_factor = self.config.init_cutoff_factor
+        elif self.config.init_fn == InitFnType.mitchell:
+            std = 1 / math.sqrt(self.config.d_model)
+            cutoff_factor = self.config.init_cutoff_factor or 3.0
+        elif self.config.init_fn == InitFnType.full_megatron:
+            std = self.config.init_std
+            cutoff_factor = self.config.init_cutoff_factor or 3.0
+        else:
+            raise NotImplementedError(self.config.init_fn)
+
+        if not self.qargs.use_quantize_model:
+            init_normal(self.att_proj, std, cutoff_factor)
+            init_normal(self.ff_proj, std, cutoff_factor)
+        else:
+            init_normal(self.BeforeAttention.att_proj, std, cutoff_factor)
+            init_normal(self.MLPResidual.ff_proj, std, cutoff_factor)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        qx: torch.Tensor,
+        sx: torch.Tensor,
+        attention_bias: torch.Tensor | None = None,
+        layer_past: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+        max_doc_len: int | None = None,
+        cu_doc_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        # Get query, key, value projections.
+        # shape:
+        #  - for regular attn q, k, v: (batch_size, seq_len, d_model)
+        #  - for multi-query attn q: (batch_size, seq_len, d_model)
+        #                      k, v: (batch_size, seq_len, d_model // n_heads)
+        #  - for group query attn q: (batch_size, seq_len, d_model)
+        #                      k, v: (batch_size, seq_len, d_model // n_kv_heads)
+
+        # import IPython
+        # IPython.embed()
+
+        if self.qargs.use_quantize_model:
+            # if False:
+            x, qkv = self.BeforeAttention(x, qx, sx)
+        else:
+            # apply norm before
+            h = self.attn_norm(x)
+
+            qkv = self.BeforeAttention.att_proj(h)
+
+        if self.config.clip_qkv is not None:
+            qkv.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
+
+        q, k, v = qkv.split(self.fused_dims, dim=-1)
+
+        # Get attention scores.
+        att, cache = self.attention(
+            q,
+            k,
+            v,
+            attention_bias,
+            layer_past=layer_past,
+            use_cache=use_cache,
+            max_doc_len=max_doc_len,
+            cu_doc_lens=cu_doc_lens,
+        )
+
+        # import IPython
+        # IPython.embed()
+        if self.qargs.use_quantize_model:
+            # if False:
+            x, qx, sx = self.AfterAttention(x, att)
+        else:
+            att = self.AfterAttention.attn_out(att)
+
+            # Add attention scores.
+            # shape: (B, T, C)
+            x = x + self.dropout(att)
+
+        if self.qargs.use_quantize_model:
+            # if False:
+            x, qx, sx = self.MLPResidual(x, qx, sx)
+        else:
+            # Add feed-forward projection.
+            # shape: (batch_size, seq_len, d_model)
+            og_x = x
+
+            x = self.ff_norm(x)
+
+            x = self.MLPResidual.ff_proj(x)
+
+            if self._activation_checkpoint_fn is not None:
+                x = self._activation_checkpoint_fn(self.act, x)  # type: ignore
+            else:
+                x = self.act(x)
+            x = self.MLPResidual.ff_out(x)
+
+            x = self.dropout(x)
+            x = og_x + x
+
+        # import IPython
+        # IPython.embed()
+
+        return x, qx, sx, cache
+
+
+class CoatOLMoLlamaBlock(OLMoBlock):
+    """
+    This is a transformer block where the output is computed as ``MLP(LN(x + Attention(LN(x))))``
+    (plus another skip connection). This block is similar to `OLMoSequentialBlock`
+    but some operations have slightly different implementations to imitate the
+    behavior of Llama.
+    """
+
+    def __init__(self, layer_id: int, config: ModelConfig, qargs: QuantActivationConfig, cache: BufferCache):
+        super().__init__(layer_id, config, qargs, cache)
+        # Layer norms.
+        self.attn_norm = LayerNorm.build(config)
+        self.ff_norm = LayerNorm.build(config)
+        self.__cache = cache
+
+        # Attention input projection. Projects x -> (q, k, v)
+        if config.multi_query_attention:
+            q_proj_out_dim = config.d_model
+            k_proj_out_dim = config.d_model // config.n_heads
+            v_proj_out_dim = config.d_model // config.n_heads
+        else:
+            q_proj_out_dim = config.d_model
+            k_proj_out_dim = config.d_model
+            v_proj_out_dim = config.d_model
+        self.q_proj = nn.Linear(config.d_model, q_proj_out_dim, bias=config.include_bias, device=config.init_device)
+        self.k_proj = nn.Linear(config.d_model, k_proj_out_dim, bias=config.include_bias, device=config.init_device)
+        self.v_proj = nn.Linear(config.d_model, v_proj_out_dim, bias=config.include_bias, device=config.init_device)
+
+        # Feed-forward input projection.
+        self.ff_proj = nn.Linear(config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device)
+
+    def reset_parameters(self):
+        super().reset_parameters()
+        self.attn_norm.reset_parameters()
+        self.ff_norm.reset_parameters()
+        # NOTE: the standard deviation for these weights does not depend on the layer.
+
+        if self.config.init_fn == InitFnType.normal:
+            std = self.config.init_std
+            cutoff_factor = self.config.init_cutoff_factor
+        elif self.config.init_fn == InitFnType.mitchell:
+            std = 1 / math.sqrt(self.config.d_model)
+            cutoff_factor = self.config.init_cutoff_factor or 3.0
+        elif self.config.init_fn == InitFnType.full_megatron:
+            std = self.config.init_std
+            cutoff_factor = self.config.init_cutoff_factor or 3.0
+        else:
+            raise NotImplementedError(self.config.init_fn)
+
+        init_normal(self.q_proj, std, cutoff_factor)
+        init_normal(self.k_proj, std, cutoff_factor)
+        init_normal(self.v_proj, std, cutoff_factor)
+        init_normal(self.ff_proj, std, cutoff_factor)
+
+    def _scaled_dot_product_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        max_doc_len: int | None = None,
+        cu_doc_lens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if max_doc_len is not None or cu_doc_lens is not None:
+            raise NotImplementedError(f"attention document masking is not implemented for {self.__class__.__name__}")
+
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
+
+        if is_causal:
+            assert attn_mask is None
+
+            query_len, key_len = q.shape[-2], k.shape[-2]  # could be different if layer_past not None
+            attn_bias = get_causal_attention_bias(self.__cache, key_len, q.device)[:, :, :query_len, :key_len]
+        elif attn_mask is not None:
+            attn_bias = attn_mask.to(q.dtype)
+        else:
+            attn_bias = torch.zeros_like(attn_weights)
+
+        attn_weights += attn_bias
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1).to(q.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=dropout_p)
+        return torch.matmul(attn_weights, v)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        qx: torch.Tensor,
+        sx: torch.Tensor,
+        attention_bias: torch.Tensor | None = None,
+        layer_past: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+        max_doc_len: int | None = None,
+        cu_doc_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        # Get query, key, value projections.
+        # shape:
+        #  - for regular attn q, k, v: (batch_size, seq_len, d_model)
+        #  - for multi-query attn q: (batch_size, seq_len, d_model)
+        #                      k, v: (batch_size, seq_len, d_model // n_heads)
+        x_normed = self.attn_norm(x)
+        q = self.q_proj(x_normed)
+        k = self.k_proj(x_normed)
+        v = self.v_proj(x_normed)
+
+        if self.config.clip_qkv is not None:
+            q.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
+            k.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
+            v.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
+
+        # Get attention scores.
+        att, cache = self.attention(
+            q,
+            k,
+            v,
+            attention_bias,
+            layer_past=layer_past,
+            use_cache=use_cache,
+            max_doc_len=max_doc_len,
+            cu_doc_lens=cu_doc_lens,
+        )
+
+        att = self.attn_out(att)  # NOTE: we move the attn_out outside the self.attention module
+
+        # Add attention scores.
+        # shape: (B, T, C)
+        x = x + self.dropout(att)
+
+        # Add feed-forward projection.
+        # shape: (batch_size, seq_len, d_model)
+        og_x = x
+        if self._activation_checkpoint_fn is not None:
+            x = self._activation_checkpoint_fn(self.ff_norm, x)  # type: ignore
+        else:
+            x = self.ff_norm(x)
+        x = self.ff_proj(x)
+        if self._activation_checkpoint_fn is not None:
+            x = self._activation_checkpoint_fn(self.act, x)  # type: ignore
+        else:
+            x = self.act(x)
+        x = self.ff_out(x)
+        x = self.dropout(x)
+        x = og_x + x
+
+        return x, cache
+
+
+class CoatOLMoBlockGroup(nn.ModuleList):
+    def __init__(self, config: ModelConfig, layer_offset: int, modules: Iterable[nn.Module] | None = None):
+        super().__init__(modules)
+        self.config = config
+        self.layer_offset = layer_offset
+        self.activation_checkpointing_strategy: ActivationCheckpointingStrategy | None = None
+        self._activation_checkpoint_fn = activation_checkpoint_function(self.config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_bias: torch.FloatTensor | None = None,
+        layers_past: list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+        use_cache: bool = False,
+        max_doc_len: int | None = None,
+        cu_doc_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor]] | None]:
+        attn_key_values: list[tuple[torch.Tensor, torch.Tensor]] | None = [] if use_cache else None
+        for block_idx, block in enumerate(self):
+            layer_past = None if layers_past is None else layers_past[block_idx]
+            block_idx += self.layer_offset
+            if should_checkpoint_block(self.activation_checkpointing_strategy, block_idx):
+                # shape: (batch_size, seq_len, d_model)
+                x, cache = self._activation_checkpoint_fn(  # type: ignore
+                    block,
+                    x,
+                    attention_bias=attention_bias,
+                    layer_past=layer_past,
+                    use_cache=use_cache,
+                    max_doc_len=max_doc_len,
+                    cu_doc_lens=cu_doc_lens,
+                )
+            else:
+                # shape: (batch_size, seq_len, d_model)
+                x, cache = block(
+                    x,
+                    attention_bias=attention_bias,
+                    layer_past=layer_past,
+                    use_cache=use_cache,
+                    max_doc_len=max_doc_len,
+                    cu_doc_lens=cu_doc_lens,
+                )
+            if attn_key_values is not None:
+                assert cache is not None
+                attn_key_values.append(cache)
+        return x, attn_key_values
+
+    def reset_parameters(self):
+        for block in self:
+            block.reset_parameters()
+
+    def set_activation_checkpointing(
+        self, strategy: ActivationCheckpointingStrategy | None, checkpoint_func: Callable | None = None
+    ):
+        self.activation_checkpointing_strategy = strategy
+        for block in self:
+            block.set_activation_checkpointing(strategy, checkpoint_func=checkpoint_func)
+
+
+class CoatOLMo(nn.Module):
+    def __init__(self, config: ModelConfig, qargs: QuantActivationConfig, init_params: bool = True):
+        super().__init__()
+        self.config = config
+        self.qargs = qargs
+        self.__cache = BufferCache()
+
+        # Validate config.
+        if self.config.alibi and self.config.flash_attention:
+            raise OLMoConfigurationError("ALiBi is currently not supported with FlashAttention")
+
+        if self.config.alibi and self.config.rope:
+            raise OLMoConfigurationError("ALiBi and RoPE are mutually exclusive")
+
+        if self.config.embedding_size is not None and self.config.embedding_size != self.config.vocab_size:
+            if self.config.embedding_size < self.config.vocab_size:
+                raise OLMoConfigurationError("embedding size should be at least as big as vocab size")
+            elif self.config.embedding_size % 128 != 0:
+                import warnings
+
+                warnings.warn(
+                    "Embedding size is not a multiple of 128! This could hurt throughput performance.", UserWarning
+                )
+
+        self.activation_checkpointing_strategy: ActivationCheckpointingStrategy | None = None
+        self._activation_checkpoint_fn: Callable = activation_checkpoint_function(self.config)
+
+        if not (
+            0 < self.config.block_group_size <= self.config.n_layers
+            and self.config.n_layers % self.config.block_group_size == 0
+        ):
+            raise OLMoConfigurationError("n layers must be divisible by block group size")
+
+        torch.backends.cuda.enable_flash_sdp(True)
+        torch.backends.cuda.enable_mem_efficient_sdp(False)  # this is super slow so make sure torch won't use it
+
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Embedding(config.embedding_size or config.vocab_size, config.d_model, device=config.init_device),
+                emb_drop=Dropout(config.embedding_dropout),
+                ln_f=LayerNorm.build(config),
+            )
+        )
+
+        blocks = [CoatOLMoBlock.build(i, config, qargs, self.__cache) for i in range(config.n_layers)]
+        if self.config.block_group_size > 1:
+            block_groups = [
+                CoatOLMoBlockGroup(config, i, blocks[i : i + config.block_group_size])
+                for i in range(0, config.n_layers, config.block_group_size)
+            ]
+            self.transformer.update({"block_groups": nn.ModuleList(block_groups)})
+        else:
+            self.transformer.update({"blocks": nn.ModuleList(blocks)})
+
+        if not (self.config.alibi or self.config.rope):
+            self.transformer.update(
+                {"wpe": nn.Embedding(config.max_sequence_length, config.d_model, device=config.init_device)}
+            )
+        if not config.weight_tying:
+            self.transformer.update(
+                {
+                    "ff_out": nn.Linear(
+                        config.d_model,
+                        config.embedding_size or config.vocab_size,
+                        bias=config.include_bias,
+                        device=config.init_device,
+                    )
+                }
+            )
+        if config.embedding_layer_norm:
+            self.transformer.update({"emb_norm": LayerNorm.build(config)})
+
+        # When `init_device="meta"` FSDP will call `reset_parameters()` to initialize weights.
+        if init_params and self.config.init_device != "meta":
+            self.reset_parameters()
+        self.__num_fwd_flops: int | None = None
+        self.__num_bck_flops: int | None = None
+
+        # Warm up cache.
+        if self.config.alibi:
+            get_causal_attention_bias(self.__cache, config.max_sequence_length, _non_meta_init_device(config))
+            self.get_alibi_attention_bias(config.max_sequence_length, _non_meta_init_device(config))
+
+        # Quantize
+        self.quantize_input_before_block = Coat_quantize_bgn(qargs)
+        self.quantize_output_after_block = Coat_quantize_end(qargs)
+
+    set_activation_checkpointing = OLMo.set_activation_checkpointing
+    device = OLMo.device
+    reset_parameters = OLMo.reset_parameters
+    get_alibi_attention_bias = OLMo.get_alibi_attention_bias
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        input_embeddings: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        attention_bias: torch.Tensor | None = None,
+        past_key_values: Sequence[tuple[torch.Tensor, torch.Tensor]] | None = None,
+        use_cache: bool = False,
+        last_logits_only: bool = False,
+        output_hidden_states: bool | None = None,
+        doc_lens: torch.Tensor | None = None,
+        max_doc_lens: Sequence[int] | None = None,
+    ) -> OLMoOutput:
+        """
+        :param input_ids: A tensor of shape `(batch_size, seq_len)`.
+        :param input_embeddings: A tensor of shape `(batch_size, seq_len, d_model)` with input
+            embeddings. When provided, it is treated as the output of the input embedding layer.
+        :param attention_mask: A tensor of shape `(batch_size, seq_len)` that indicates
+            which input IDs are masked. A `1` value in the mask means that
+            the corresponding input ID should *not* be ignored. A `0` means
+            that the corresponding input ID is masked.
+
+            This has the same meaning as the `attention_mask` in HuggingFace's `transformers`
+            library.
+        :param attention_bias: A tensor of shape `(batch_size, 1, seq_len, seq_len)`,
+            `(1, 1, seq_len, seq_len)`, or `(seq_len, seq_len)`. This is used
+            to introduce causal or other biases.
+
+            If the tensor is a bool or byte tensor, a `True` or `1` at `attention_bias[:, :, i, j]`
+            indicates that the i-th element in the sequence is allowed to attend to the j-th
+            element in the sequence.
+
+            If the tensor is a float tensor, it will just be added to the attention
+            scores before the softmax.
+
+            The default is causal, which corresponds to a lower-diagonal byte matrix of ones.
+        :param past_key_values: Pre-computed keys and values for each attention block.
+            Can be used to speed up sequential decoding. The `input_ids` which have
+            their past given to this model should not be passed as `input_ids` as they have already been computed.
+        :param use_cache: If `True`, return key and value tensors for each block.
+        :param last_logits_only: If `True`, only compute the logits for the last token of each sequence.
+            This can speed up decoding when you only care about the next token.
+        :param doc_lens: Document lengths to use in attention for intra-document masking.
+            Shape `(batch_size, max_docs)`.
+        :param max_doc_lens: Maximum document length for each instance in the batch.
+        """
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+
+        if past_key_values:
+            assert len(past_key_values) == self.config.n_layers
+
+        batch_size, seq_len = input_ids.size() if input_embeddings is None else input_embeddings.size()[:2]
+        if past_key_values is None:
+            past_length = 0
+        else:
+            past_length = past_key_values[0][0].size(-2)
+
+        max_doc_len: int | None = None
+        cu_doc_lens: torch.Tensor | None = None
+        if doc_lens is not None and max_doc_lens is not None:
+            max_doc_len = max(max_doc_lens)
+            cu_doc_lens = get_cumulative_document_lengths(doc_lens)
+
+        # Get embeddings of input.
+        # shape: (batch_size, seq_len, d_model)
+        x = self.transformer.wte(input_ids) if input_embeddings is None else input_embeddings  # type: ignore
+
+        # Apply embedding layer norm.
+        if self.config.embedding_layer_norm:
+            x = self.transformer.emb_norm(x)
+
+        if not (self.config.alibi or self.config.rope):
+            # Get positional embeddings.
+            # shape: (1, seq_len)
+            pos = torch.arange(past_length, past_length + seq_len, dtype=torch.long, device=x.device).unsqueeze(0)
+            # shape: (1, seq_len, d_model)
+            pos_emb = self.transformer.wpe(pos)  # type: ignore
+            x = pos_emb + x
+
+        # Apply dropout.
+        # shape: (batch_size, seq_len, d_model)
+        x = self.transformer.emb_drop(x)  # type: ignore
+
+        # Transform the attention mask into what the blocks expect.
+        if attention_mask is not None:
+            # shape: (batch_size, 1, 1, seq_len)
+            attention_mask = attention_mask.to(dtype=torch.float).view(batch_size, -1)[:, None, None, :]
+            attention_mask = (1.0 - attention_mask) * torch.finfo(attention_mask.dtype).min
+
+        # Merge attention mask with attention bias.
+        if (
+            attention_bias is not None
+            or attention_mask is not None
+            or self.config.alibi
+            # NOTE (epwalsh): we need to initialize the attn bias in order for attn to work properly
+            # with key+value cache. Otherwise `F.scaled_dot_product_attention()` doesn't seem to compute
+            # scores correctly.
+            or past_key_values is not None
+        ):
+            if attention_bias is None and self.config.alibi:
+                attention_bias = get_causal_attention_bias(
+                    self.__cache, past_length + seq_len, x.device
+                ) + self.get_alibi_attention_bias(past_length + seq_len, x.device)
+            elif attention_bias is None:
+                attention_bias = get_causal_attention_bias(self.__cache, past_length + seq_len, x.device)
+            elif attention_bias.dtype in (torch.int8, torch.bool):
+                attention_bias = attention_bias.to(dtype=torch.float)
+                attention_bias.masked_fill_(attention_bias == 0.0, torch.finfo(attention_bias.dtype).min)
+
+            # Transform to the right shape and data type.
+            mask_len = seq_len
+            if attention_mask is not None:
+                mask_len = attention_mask.shape[-1]
+            elif past_key_values is not None:
+                mask_len = past_key_values[0][0].shape[-2] + seq_len
+            attention_bias = attention_bias[:, :, :mask_len, :mask_len].to(dtype=torch.float)
+
+            # Add in the masking bias.
+            if attention_mask is not None:
+                attention_bias = attention_bias + attention_mask
+                # Might get -infs after adding attention mask, since dtype.min + dtype.min = -inf.
+                # `F.scaled_dot_product_attention()` doesn't handle -inf like you'd expect, instead
+                # it can produce NaNs.
+                ensure_finite_(attention_bias, check_neg_inf=True, check_pos_inf=False)
+
+        attn_key_values: list[tuple[torch.Tensor, torch.Tensor]] | None = [] if use_cache else None
+
+        # decoder layers
+        all_hidden_states = []
+
+        # Prepare the input for COAT decoderlayer
+        x, qx, sx = self.quantize_input_before_block(x)
+
+        # Apply blocks one-by-one.
+        if self.config.block_group_size == 1:
+            for block_idx, block in enumerate(self.transformer.blocks):
+                if output_hidden_states:
+                    # add hidden states
+                    all_hidden_states.append(x)
+
+                layer_past = None if past_key_values is None else past_key_values[block_idx]
+                if should_checkpoint_block(self.activation_checkpointing_strategy, block_idx):
+                    # shape: (batch_size, seq_len, d_model)
+                    x, qx, sx, cache = self._activation_checkpoint_fn(
+                        block,
+                        x,
+                        qx,
+                        sx,
+                        attention_bias=attention_bias,
+                        layer_past=layer_past,
+                        use_cache=use_cache,
+                        max_doc_len=max_doc_len,
+                        cu_doc_lens=cu_doc_lens,
+                    )
+                else:
+                    # shape: (batch_size, seq_len, d_model)
+                    x, qx, sx, cache = block(
+                        x,
+                        qx,
+                        sx,
+                        attention_bias=attention_bias,
+                        layer_past=layer_past,
+                        use_cache=use_cache,
+                        max_doc_len=max_doc_len,
+                        cu_doc_lens=cu_doc_lens,
+                    )
+
+                if attn_key_values is not None:
+                    assert cache is not None
+                    attn_key_values.append(cache)
+        else:
+            for group_idx, block_group in enumerate(self.transformer.block_groups):
+                if output_hidden_states:
+                    # add hidden states
+                    all_hidden_states.append(x)
+
+                layers_past = (
+                    None
+                    if past_key_values is None
+                    else past_key_values[
+                        group_idx * self.config.block_group_size : (group_idx + 1) * self.config.block_group_size
+                    ]
+                )
+                x, cache = block_group(
+                    x,
+                    attention_bias=attention_bias,
+                    layers_past=layers_past,
+                    use_cache=use_cache,
+                    max_doc_len=max_doc_len,
+                    cu_doc_lens=cu_doc_lens,
+                )
+                if attn_key_values is not None:
+                    assert cache is not None
+                    attn_key_values.extend(cache)
+
+        # Summarize the output of the Decoder Layer
+        x = self.quantize_output_after_block(x, qx, sx)
+
+        if last_logits_only:
+            # shape: (batch_size, 1, d_model)
+            x = x[:, -1, :].unsqueeze(1)
+
+        # Apply final layer norm.
+        # shape: (batch_size, seq_len or 1, d_model)
+        x = self.transformer.ln_f(x)  # type: ignore
+        if output_hidden_states:
+            # add final hidden state post-final-layernorm, following HuggingFace's convention
+            all_hidden_states.append(x)
+
+        # Get logits.
+        # shape: (batch_size, seq_len or 1, vocab_size)
+        if self.config.weight_tying:
+            logits = F.linear(x, self.transformer.wte.weight, None)  # type: ignore
+        else:
+            logits = self.transformer.ff_out(x)  # type: ignore
+        if self.config.scale_logits:
+            logits.mul_(1 / math.sqrt(self.config.d_model))
+
+        return OLMoOutput(
+            logits=logits,
+            attn_key_values=attn_key_values,
+            hidden_states=tuple(all_hidden_states) if output_hidden_states else None,
+        )
+
+    def get_fsdp_wrap_policy(self, wrap_strategy: FSDPWrapStrategy | None = None):
+        if wrap_strategy is None:
+            return None
+
+        # The 'recurse' mode for the wrap function does not behave like you'd expect.
+        # Even if we return False, it may still recurse because PyTorch does what it wants,
+        # not what you want. This causes issues when, for example, we want to wrap 'ff_out' (a linear layer)
+        # but not other linear layers within a block.
+        # So we have to explicitly tell PyTorch which linear layers to wrap, and we also just
+        # return True in 'recurse' mode for simplicity.
+        size_based_module_to_wrap = {self.transformer.wte}
+        if hasattr(self.transformer, "ff_out"):
+            size_based_module_to_wrap.add(self.transformer.ff_out)
+
+        if wrap_strategy == FSDPWrapStrategy.by_block:
+
+            def fsdp_wrap_fn(module, recurse: bool = True, nonwrapped_numel: int = 0):
+                del nonwrapped_numel
+                wrap = isinstance(module, CoatOLMoBlock)
+                if recurse:
+                    return True
+                else:
+                    return wrap
+
+            return fsdp_wrap_fn
+        elif wrap_strategy == FSDPWrapStrategy.by_block_and_size:
+
+            def fsdp_wrap_fn(module, recurse: bool = True, nonwrapped_numel: int = 0):
+                del nonwrapped_numel
+                wrap = isinstance(module, (CoatOLMoBlock,)) or module in size_based_module_to_wrap
+                if recurse:
+                    return True
+                else:
+                    return wrap
+
+            return fsdp_wrap_fn
+        elif wrap_strategy == FSDPWrapStrategy.by_block_group:
+            if self.config.block_group_size <= 1:
+                raise OLMoConfigurationError(
+                    "'by_block_group' FSDP wrapping strategy requires block group size greater than 1"
+                )
+
+            def fsdp_wrap_fn(module, recurse: bool = True, nonwrapped_numel: int = 0):
+                del nonwrapped_numel
+                wrap = isinstance(module, CoatOLMoBlockGroup)
+                if recurse:
+                    return True
+                else:
+                    return wrap
+
+            return fsdp_wrap_fn
+        elif wrap_strategy == FSDPWrapStrategy.by_block_group_and_size:
+            if self.config.block_group_size <= 1:
+                raise OLMoConfigurationError(
+                    "'by_block_group_and_size' FSDP wrapping strategy requires block group size greater than 1"
+                )
+
+            def fsdp_wrap_fn(module, recurse: bool = True, nonwrapped_numel: int = 0):
+                del nonwrapped_numel
+                wrap = isinstance(module, (CoatOLMoBlockGroup,)) or module in size_based_module_to_wrap
+                if recurse:
+                    return True
+                else:
+                    return wrap
+
+            return fsdp_wrap_fn
+        elif wrap_strategy == FSDPWrapStrategy.size_based:
+            from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
+
+            return size_based_auto_wrap_policy
+        elif wrap_strategy in {
+            FSDPWrapStrategy.one_in_two,
+            FSDPWrapStrategy.one_in_three,
+            FSDPWrapStrategy.one_in_four,
+            FSDPWrapStrategy.one_in_five,
+        }:
+            c = {
+                FSDPWrapStrategy.one_in_two: 2,
+                FSDPWrapStrategy.one_in_three: 3,
+                FSDPWrapStrategy.one_in_four: 4,
+                FSDPWrapStrategy.one_in_five: 5,
+            }[wrap_strategy]
+
+            def fsdp_wrap_fn(module, recurse: bool = True, nonwrapped_numel: int = 0):
+                del nonwrapped_numel
+                wrap = isinstance(module, CoatOLMoBlock) and module.layer_id % c == 0
+                if recurse:
+                    return True
+                else:
+                    return wrap
+
+            return fsdp_wrap_fn
+        else:
+            raise NotImplementedError(wrap_strategy)
+
+    num_params = OLMo.num_params
+
+    @property
+    def num_fwd_flops(self):
+        if self.__num_fwd_flops:
+            return self.__num_fwd_flops
+
+        # embedding table is just a lookup in the forward pass
+        n_params = self.num_params(include_embedding=False)
+        # the number of parameters is approximately the number of multiply-accumulates (MAC) in the network
+        # each MAC has 2 FLOPs - we multiply by 2 ie 2 * n_param
+        # this gets us FLOPs / token
+        params_flops_per_token = 2 * n_params
+        # there are 2 FLOPS per mac; there is A=Q*K^T and out=A*V ops (ie mult by 2)
+        attn_flops_per_token = self.config.n_layers * 2 * 2 * (self.config.d_model * self.config.max_sequence_length)
+        self.__num_fwd_flops = params_flops_per_token + attn_flops_per_token
+        return self.__num_fwd_flops
+
+    @property
+    def num_bck_flops(self):
+        if self.__num_bck_flops:
+            return self.__num_bck_flops
+
+        n_params = self.num_params()
+        params_flops_per_token = 4 * n_params
+        attn_flops_per_token = self.config.n_layers * 8 * (self.config.d_model * self.config.max_sequence_length)
+        self.__num_bck_flops = params_flops_per_token + attn_flops_per_token
+        return self.__num_bck_flops
+
+    generate = OLMo.generate
+
+    @classmethod
+    def from_checkpoint(
+        cls, checkpoint_dir: PathOrStr, device: str = "cpu", checkpoint_type: CheckpointType | None = None
+    ) -> CoatOLMo:
+        """
+        Load an OLMo model from a checkpoint.
+        """
+        from olmo.util import resource_path
+
+        # Guess checkpoint type.
+        if checkpoint_type is None:
+            try:
+                if resource_path(checkpoint_dir, "model.pt").is_file():
+                    checkpoint_type = CheckpointType.unsharded
+                else:
+                    checkpoint_type = CheckpointType.sharded
+            except FileNotFoundError:
+                checkpoint_type = CheckpointType.sharded
+
+        # Load config.
+        config_path = resource_path(checkpoint_dir, "config.yaml")
+        model_config = ModelConfig.load(config_path, key="model", validate_paths=False)
+
+        if checkpoint_type == CheckpointType.unsharded:
+            # Initialize model (always on CPU to start with so we don't run out of GPU memory).
+            model_config.init_device = "cpu"
+            model = CoatOLMo(model_config)
+
+            # Load state dict directly to target device.
+            state_dict_path = resource_path(checkpoint_dir, "model.pt")
+            state_dict = torch.load(state_dict_path, map_location="cpu")
+            model.load_state_dict(model._make_state_dict_compatible(state_dict)[0])
+            model = model.to(torch.device(device))
+        else:
+            train_config = TrainConfig.load(config_path)
+            if train_config.sharded_checkpointer == ShardedCheckpointerType.olmo_core:
+                from olmo_core.distributed.checkpoint import load_model_and_optim_state  # type: ignore
+
+                model_config.init_device = device
+                model = CoatOLMo(model_config)
+                load_model_and_optim_state(checkpoint_dir, model)
+            else:
+                # train_config.sharded_checkpointer == ShardedCheckpointerType.torch_new
+                from olmo.checkpoint import load_model_state
+
+                # Initialize model on target device. In this case the state dict is loaded in-place
+                # so it's not necessary to start on CPU if the target device is a GPU.
+                model_config.init_device = device
+                model = CoatOLMo(model_config)
+
+                # Load state dict in place.
+                load_model_state(checkpoint_dir, model)
+
+        return model.eval()
+
+    def _make_state_dict_compatible(
+        self, state_dict: dict[str, torch.Tensor]
+    ) -> tuple[dict[str, torch.Tensor], dict[str, set[str]]]:
+        """
+        Handles some cases where the state dict is valid yet may need to be transformed in order to
+        be loaded.
+
+        This modifies the state dict in-place and also returns it, along with a mapping of original key
+        names to new key names in cases where the keys were simply renamed. That mapping can be used
+        to make a corresponding optimizer state dict compatible as well.
+        """
+        import re
+        from fnmatch import fnmatch
+
+        new_keys_to_og_keys: dict[str, str] = {}
+
+        # Remove "_fsdp_wrapped_module." prefix from all keys. We don't want this prefix when the model is
+        # not wrapped in FSDP. And when the model is wrapped in FSDP, loading this state dict will still work
+        # fine without the prefixes. This also simplifies the other steps below.
+        for key in list(state_dict.keys()):
+            state_dict[(new_key := key.replace("_fsdp_wrapped_module.", ""))] = state_dict.pop(key)
+            new_keys_to_og_keys[new_key] = key
+
+        # For backwards compatibility prior to fixing https://github.com/allenai/LLM/issues/222
+        if self.config.block_type == BlockType.sequential:
+            for key in list(state_dict.keys()):
+                if fnmatch(key, "transformer.*.norm.weight"):
+                    tensor = state_dict.pop(key)
+                    state_dict[(new_key := key.replace("norm.weight", "attn_norm.weight"))] = tensor
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    state_dict[(new_key := key.replace("norm.weight", "ff_norm.weight"))] = tensor.clone()
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    del new_keys_to_og_keys[key]
+                elif fnmatch(key, "transformer.*.norm.bias"):
+                    tensor = state_dict.pop(key)
+                    state_dict[(new_key := key.replace("norm.bias", "attn_norm.bias"))] = tensor
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    state_dict[(new_key := key.replace("norm.bias", "ff_norm.bias"))] = tensor.clone()
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    del new_keys_to_og_keys[key]
+
+        # Realquantization will change the place the linear layers happen
+        if self.qargs.use_quantize_model == "coat_real":
+            for key in list(state_dict.keys()):
+                if fnmatch(key, "transformer.blocks.*.att_proj.weight") and "BeforeAttention" not in key:
+                    tensor = state_dict.pop(key)
+                    state_dict[(new_key := key.replace("att_proj.weight", "BeforeAttention.att_proj.weight"))] = tensor
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    del new_keys_to_og_keys[key]
+                elif fnmatch(key, "transformer.blocks.*.attn_out.weight") and "AfterAttention" not in key:
+                    tensor = state_dict.pop(key)
+                    state_dict[(new_key := key.replace("attn_out.weight", "AfterAttention.attn_out.weight"))] = tensor
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    del new_keys_to_og_keys[key]
+                elif fnmatch(key, "transformer.blocks.*.ff_proj.weight") and "MLPResidual" not in key:
+                    tensor = state_dict.pop(key)
+                    state_dict[(new_key := key.replace("ff_proj.weight", "MLPResidual.ff_proj.weight"))] = tensor
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    del new_keys_to_og_keys[key]
+                elif fnmatch(key, "transformer.blocks.*.ff_out.weight") and "MLPResidual" not in key:
+                    tensor = state_dict.pop(key)
+                    state_dict[(new_key := key.replace("ff_out.weight", "MLPResidual.ff_out.weight"))] = tensor
+                    new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key]
+                    del new_keys_to_og_keys[key]
+
+        # For loading a state dict that was saved with a different `block_group_size`.
+        if "transformer.block_groups.0.0.attn_out.weight" in state_dict.keys():
+            state_dict_block_group_size = len(
+                [k for k in state_dict.keys() if fnmatch(k, "transformer.block_groups.0.*.attn_out.weight")]
+            )
+        else:
+            state_dict_block_group_size = 1
+        if self.config.block_group_size != state_dict_block_group_size:
+            log.info(
+                f"Regrouping state dict blocks from group size {state_dict_block_group_size} to "
+                f"group size {self.config.block_group_size}"
+            )
+            # For simplicity we're first going to flatten out the block groups in the state dict (if necessary)
+            # and then (re-)group them into the right block sizes.
+            if state_dict_block_group_size > 1:
+                for key in list(state_dict.keys()):
+                    if (m := re.match(r"transformer.block_groups\.(\d+)\.(\d+)\..*", key)) is not None:
+                        group_idx, group_block_idx = int(m.group(1)), int(m.group(2))
+                        block_idx = (group_idx * state_dict_block_group_size) + group_block_idx
+                        state_dict[
+                            (
+                                new_key := key.replace(
+                                    f"block_groups.{group_idx}.{group_block_idx}.", f"blocks.{block_idx}."
+                                )
+                            )
+                        ] = state_dict.pop(key)
+                        new_keys_to_og_keys[new_key] = new_keys_to_og_keys.pop(key)
+
+            if self.config.block_group_size > 1:
+                # Group the state dict blocks into the right block size.
+                for key in list(state_dict.keys()):
+                    if (m := re.match(r"transformer.blocks\.(\d+)\..*", key)) is not None:
+                        block_idx = int(m.group(1))
+                        group_idx, group_block_idx = (
+                            block_idx // self.config.block_group_size,
+                            block_idx % self.config.block_group_size,
+                        )
+                        state_dict[
+                            (
+                                new_key := key.replace(
+                                    f"blocks.{block_idx}.", f"block_groups.{group_idx}.{group_block_idx}."
+                                )
+                            )
+                        ] = state_dict.pop(key)
+                        new_keys_to_og_keys[new_key] = new_keys_to_og_keys.pop(key)
+
+        og_keys_to_new: dict[str, set[str]] = defaultdict(set)
+        for new_key, og_key in new_keys_to_og_keys.items():
+            og_keys_to_new[og_key].add(new_key)
+
+        return state_dict, og_keys_to_new

llava/model/coat/activation/real_quantization/__init__.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Activation
+# Utils
+from ._dequantize import fp8_dequantize
+from ._division import fp8_division
+from ._division_transpose import fp8_division_transpose
+from ._quantize import fp8_quantize
+from ._quantize_pertensor import fp8_quantize_pertensor
+from ._quantize_pertensor_transpose import fp8_quantize_pertensor_transpose
+from ._transpose import fp8_transpose
+from .add_bwd import fp8_add_Ifp_Ifp_Ofp_Opt
+from .add_fwd import fp8_add_Ifp_Ifp_Ofp_Og16
+
+# Normalization
+from .func_layernorm_noparam import fp8_layernorm_noparam_backward, fp8_layernorm_noparam_forward
+from .func_quantize import Coat_quantize_bgn, Coat_quantize_end
+from .func_rmsnorm import fp8_rmsnorm_backward, fp8_rmsnorm_forward
+from .gelu_bwd import fp8_gelu_backward
+from .gelu_fwd import fp8_gelu_forward
+
+# linear and add
+from .linear import fp8_linear_backward, fp8_linear_forward
+from .mul_bwd import fp8_mul_backward
+from .mul_fwd import fp8_mul_forward
+from .silu_bwd import fp8_silu_backward
+from .silu_fwd import fp8_silu_forward

llava/model/coat/activation/real_quantization/_dequantize.py

@@ -0,0 +1,162 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, get_configs_io_block
+
+"""Quantize Operator"""
+"""Input uses 1 * 16 group quantization"""
+"""Output uses 1 * 16 group quantization"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_dequantize_kernel(
+    output_ptr,  # output
+    input_ptr,
+    input_scale_ptr,  # input
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,  # shape
+    input_stride_0,
+    input_stride_1,  # input stride
+    s_input_stride_0,
+    s_input_stride_1,  # scale of output stride
+    output_stride_0,
+    output_stride_1,  # output stride
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+    # input ptr
+    scale_input_ptr = tl.make_block_ptr(
+        base=input_scale_ptr,
+        shape=(M, SN),
+        strides=(s_input_stride_0, s_input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_SN),
+        block_shape=(BLOCK_M, BLOCK_SN),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+    scale_input = tl.load(scale_input_ptr)
+
+    input = input.to(tl.float32)
+    scale_input = scale_input.to(tl.float32)
+
+    # Dequantize and gelu calculation
+    scale_input = tl.reshape(scale_input, (BLOCK_M, BLOCK_SN, 1))
+    input = tl.reshape(input, (BLOCK_M, BLOCK_SN, QB))
+    output = input * scale_input
+
+    output = tl.reshape(output, (BLOCK_M, BLOCK_N))
+    output = output.to(output_ptr.dtype.element_ty)
+
+    # debug
+    # gelu_output = input
+    # scale_output = scale_input
+
+    # pointers
+    output_block_ptr = tl.make_block_ptr(
+        base=output_ptr,
+        shape=(M, N),
+        strides=(output_stride_0, output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    tl.store(output_block_ptr, output, boundary_check=(0, 1))
+
+
+def fp8_dequantize(x, s_x, QB):
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x.shape) == 3:
+        batched = True
+        BS = x.shape[0]
+        x = x.reshape(-1, x.shape[-1])
+        s_x = s_x.reshape(-1, s_x.shape[-1])
+
+    # defining the input and output tensor
+    M, N = x.shape
+    SN = N // QB
+
+    y = torch.empty_like(x, dtype=torch.bfloat16)
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_dequantize_kernel[grid](
+        y,
+        x,
+        s_x,
+        M,
+        N,
+        SN,
+        QB,
+        x.stride(0),
+        x.stride(1),
+        s_x.stride(0),
+        s_x.stride(1),
+        y.stride(0),
+        y.stride(1),
+    )
+
+    # Recover 2D to 3D
+    if batched:
+        y = y.reshape(BS, -1, y.shape[-1])
+
+    return y

llava/model/coat/activation/real_quantization/_division.py

@@ -0,0 +1,212 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, convert_fp8_to_embit, convert_str_to_fp8, get_configs_io_block
+
+"""Quantize and Transpose Operator"""
+"""Input uses 1 * 16 group quantization"""
+"""Output uses 1 * 16 group quantization"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_division_kernel(
+    output_ptr,  # output
+    input_ptr,
+    input_scale_ptr,  # input
+    noise_ptr,  # noise for stochastic
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,
+    fp8_max,
+    e_bit: tl.constexpr,
+    m_bit: tl.constexpr,  # shape
+    input_stride_0,
+    input_stride_1,  # input stride
+    output_stride_0,
+    output_stride_1,  # output stride
+    SCALE_MIN_THRES: tl.constexpr,  # We do not use it since we believe SCALE_MIN_THRES should be used in previous kernel when calculating scaling factor
+    STOCHASTIC: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+    input = input.to(tl.float32)
+    scale_output = tl.load(input_scale_ptr)
+    scale_output = scale_output.to(tl.float32)
+
+    output = tl.reshape(input, (BLOCK_M, BLOCK_SN, QB))
+
+    # Quantize Scale calculation
+    # Quantize
+    output = tl.fdiv(output, scale_output)
+    output = tl.reshape(output, (BLOCK_M, BLOCK_N))
+
+    if STOCHASTIC:
+        # noise_block_ptr = tl.make_block_ptr(
+        #     base=noise_ptr,
+        #     shape=(M, N),
+        #     strides=(input_stride_0, input_stride_1),
+        #     offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        #     block_shape=(BLOCK_M, BLOCK_N),
+        #     order=(1, 0)
+        # )
+        # noise = tl.load(noise_block_ptr)
+
+        offs_m = pid_dim0 * BLOCK_M + tl.arange(0, BLOCK_M)
+        offs_n = pid_dim1 * BLOCK_N + tl.arange(0, BLOCK_N)
+        noise_offset = offs_m[:, None] * input_stride_0 + offs_n[None, :] * input_stride_1
+        noise = tl.rand(0, noise_offset)
+
+        output = _stochastic_rounding(output, noise, e_bit, m_bit)
+
+    output = output.to(output_ptr.type.element_ty)
+
+    # pointers
+    output_block_ptr = tl.make_block_ptr(
+        base=output_ptr,
+        shape=(M, N),
+        strides=(output_stride_0, output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    tl.store(output_block_ptr, output, boundary_check=(0, 1))
+
+
+@triton.jit
+def _stochastic_rounding(output, noise, e_bit: tl.constexpr, m_bit: tl.constexpr):
+    subnormal_min = tl.exp2(2 - tl.exp2(e_bit - 1) - m_bit)
+    # subnormal_should_be = tl.exp2(2 - tl.exp2(e_bit) - 1)
+
+    output_int32 = tl.cast(output, tl.int32, bitcast=True)
+    output_int32 = output_int32 & 0x7F800000
+    output_float32 = tl.cast(output_int32, tl.float32, bitcast=True)
+    output_exp = tl.maximum(output_float32, subnormal_min)
+
+    noise_rescale = tl.exp2(m_bit) + (output_exp == subnormal_min) * (
+        1 - tl.exp2(m_bit)
+    )  # 2^m_bit for normal, 1 for subnormal
+
+    noise = output_exp * noise / noise_rescale
+    sign = 1 - 2 * libdevice.signbit(output)
+    output = tl.abs(output) + noise
+
+    minmax_ratio = 2 + (output_exp == subnormal_min) * (tl.exp2(m_bit) - 2)  # 2 for normal, and 2^M for subnormal
+    output = sign * tl.clamp(output, min=output_exp, max=minmax_ratio * output_exp)
+
+    return output
+
+
+def fp8_division(x, QB, fp8type, s_y=None, stochastic=False):
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x.shape) == 3:
+        batched = True
+        BS = x.shape[0]
+        x = x.reshape(-1, x.shape[-1])
+
+    if stochastic:
+        # noise = torch.zeros_like(x, dtype=torch.float32).uniform_(-0.5, 0.5)
+        noise = None
+    else:
+        noise = None
+
+    # defining the input and output tensor
+    M, N = x.shape
+    SN = N // QB
+
+    if isinstance(fp8type, str):
+        fp8type = convert_str_to_fp8[fp8type]
+
+    y = torch.empty_like(x, dtype=fp8type)
+    fp8MaxValue = FP8_MAX_VALUE[fp8type]  # E4M3 and E5M2 have different max value
+    e_bit, m_bit = convert_fp8_to_embit[fp8type]
+
+    if s_y is None:
+        s_y = (x.abs().max() + SCALE_MIN_THRES) / fp8MaxValue
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_division_kernel[grid](
+        y,
+        x,
+        s_y,
+        noise,
+        M,
+        N,
+        SN,
+        QB,
+        fp8MaxValue,
+        e_bit,
+        m_bit,
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+        SCALE_MIN_THRES=SCALE_MIN_THRES,
+        STOCHASTIC=stochastic,
+    )
+
+    # Recover 2D to 3D
+    if batched:
+        y = y.reshape(BS, -1, y.shape[-1])
+
+    return y, s_y  # y_t is expected to be 2D tensor

llava/model/coat/activation/real_quantization/_division_transpose.py

@@ -0,0 +1,215 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from ._division import _stochastic_rounding
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, convert_fp8_to_embit, convert_str_to_fp8, get_configs_io_block
+
+"""Division and Transpose Operator"""
+"""Input uses full-precision/BF16"""
+"""Output uses per tensor quantization"""
+"""Output_t uses per tensor quantization and is transposed, but is flattened to 2D"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),  # triton.Config({'BLOCK_M': 1, 'BLOCK_N': 16}, num_stages=4, num_warps=1,)
+    # configs=[triton.Config({'BLOCK_M': 1, 'BLOCK_N': 16}, num_stages=4, num_warps=1,)], #
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_division_transpose_kernel(
+    output_ptr,
+    output_t_ptr,  # output
+    input_ptr,
+    input_scale_ptr,  # input
+    noise_ptr,  # noise for stochastic
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,
+    fp8_max,
+    e_bit,
+    m_bit,  # shape
+    input_stride_0,
+    input_stride_1,  # input stride
+    output_stride_0,
+    output_stride_1,  # output stride
+    output_t_stride_0,
+    output_t_stride_1,  # output stride
+    SCALE_MIN_THRES: tl.constexpr,  # We do not use it since we believe SCALE_MIN_THRES should be used in previous kernel when calculating scaling factor
+    STOCHASTIC: tl.constexpr,
+    ONLY_TRANSPOSED: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+    input = input.to(tl.float32)
+    scale_output = tl.load(input_scale_ptr)
+    scale_output = scale_output.to(tl.float32)
+
+    output = tl.reshape(input, (BLOCK_M, BLOCK_SN, QB))
+
+    # Quantize Scale calculation
+    # Quantize
+    output = tl.fdiv(output, scale_output)
+    output = tl.reshape(output, (BLOCK_M, BLOCK_N))
+
+    if STOCHASTIC:
+        # noise_block_ptr = tl.make_block_ptr(
+        #     base=noise_ptr,
+        #     shape=(M, N),
+        #     strides=(input_stride_0, input_stride_1),
+        #     offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        #     block_shape=(BLOCK_M, BLOCK_N),
+        #     order=(1, 0)
+        # )
+        # noise = tl.load(noise_block_ptr)
+
+        offs_m = pid_dim0 * BLOCK_M + tl.arange(0, BLOCK_M)
+        offs_n = pid_dim1 * BLOCK_N + tl.arange(0, BLOCK_N)
+        noise_offset = offs_m[:, None] * input_stride_0 + offs_n[None, :] * input_stride_1
+        noise = tl.rand(0, noise_offset)
+
+        output = _stochastic_rounding(output, noise, e_bit, m_bit)
+
+    output = output.to(output_ptr.type.element_ty)
+    # tl.device_print("3: ", output)
+    output_t = tl.trans(output)
+
+    # pointers
+    output_block_ptr = tl.make_block_ptr(
+        base=output_ptr,
+        shape=(M, N),
+        strides=(output_stride_0, output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+    output_t_block_ptr = tl.make_block_ptr(
+        base=output_t_ptr,
+        shape=(N, M),
+        strides=(output_t_stride_0, output_t_stride_1),
+        offsets=(pid_dim1 * BLOCK_N, pid_dim0 * BLOCK_M),
+        block_shape=(BLOCK_N, BLOCK_M),
+        order=(1, 0),
+    )
+    if not ONLY_TRANSPOSED:
+        tl.store(output_block_ptr, output, boundary_check=(0, 1))
+    tl.store(output_t_block_ptr, output_t, boundary_check=(0, 1))
+
+
+def fp8_division_transpose(x, QB, fp8type, s_y=None, stochastic=False, only_transposed=False):
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x.shape) == 3:
+        batched = True
+        BS = x.shape[0]
+        x = x.reshape(-1, x.shape[-1])
+
+    if stochastic:
+        # noise = torch.empty_like(x, dtype=torch.float32).uniform_(-0.5, 0.5)
+        noise = None
+    else:
+        noise = None
+
+    # defining the input and output tensor
+    M, N = x.shape
+    SN = N // QB
+
+    if isinstance(fp8type, str):
+        fp8type = convert_str_to_fp8[fp8type]
+
+    y = torch.empty_like(x, dtype=fp8type)
+    y_t = torch.empty((N, M), dtype=fp8type, device=x.device)
+    fp8MaxValue = FP8_MAX_VALUE[fp8type]  # E4M3 and E5M2 have different max value
+    e_bit, m_bit = convert_fp8_to_embit[fp8type]
+
+    if s_y is None:
+        # print("Warning: do not specify s_y in fp8_division_transpose")
+        s_y = (x.abs().max() + SCALE_MIN_THRES) / fp8MaxValue
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_division_transpose_kernel[grid](
+        y,
+        y_t,
+        x,
+        s_y,
+        noise,
+        M,
+        N,
+        SN,
+        QB,
+        fp8MaxValue,
+        e_bit,
+        m_bit,
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+        y_t.stride(0),
+        y_t.stride(1),
+        SCALE_MIN_THRES=SCALE_MIN_THRES,
+        STOCHASTIC=stochastic,
+        ONLY_TRANSPOSED=only_transposed,
+    )
+
+    if not only_transposed:
+        # Recover 2D to 3D
+        if batched:
+            y = y.reshape(BS, -1, y.shape[-1])
+
+        return y, s_y, y_t  # y_t is expected to be 2D tensor
+    else:
+        return y_t, s_y

llava/model/coat/activation/real_quantization/_memory_io.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+CONST_BLOCK = 32
+
+# The kernel with 1 load operation and 4 store operation
+def get_configs_io_block():
+    configs = []
+    for nstages in [3, 4, 5, 6]:
+        for block_m in [32, 64, 128]:
+            for block_n in [32, 64, 128]:
+                for nwarps in [4, 8, 16, 32]:
+                    configs.append(
+                        triton.Config(
+                            {"BLOCK_M": block_m, "BLOCK_N": block_n},
+                            num_stages=nstages,
+                            num_warps=nwarps,
+                        )
+                    )
+    return configs
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.jit
+def bench_memory_io_kernel_forward(
+    output_ptr,
+    input_ptr,
+    M,
+    N,
+    B: tl.constexpr,
+    input_stride_0,
+    input_stride_1,
+    output_stride_0,
+    output_stride_1,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_M = tl.cdiv(M, BLOCK_M)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+    input = input.to(tl.float32)
+
+    output = input * 2
+
+    # pointers
+    output_block_ptr = tl.make_block_ptr(
+        base=output_ptr,
+        shape=(M, N),
+        strides=(output_stride_0, output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    output = output.to(output_ptr.type.element_ty)
+    tl.store(output_block_ptr, output, boundary_check=(0, 1))
+
+
+def bench_memory_io_forward(x, B):
+    # defining the input and output tensor
+    M, N = x.shape
+
+    y = torch.empty_like(x, dtype=x.dtype)
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    bench_memory_io_kernel_forward[grid](
+        y,
+        x,
+        M,
+        N,
+        B,
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+    )
+    return y
+
+
+configs = []
+for SL in [8192]:
+    configs.append(
+        triton.testing.Benchmark(  # test different matrix size influence
+            x_names=["CDIM"],
+            x_vals=[1024, 2048, 4096, 8192],
+            line_arg="dtype",
+            line_vals=[torch.int8, torch.float16, torch.float32],
+            line_names=["float8", "float16", "float32"],
+            styles=[("blue", "-"), ("green", "-"), ("red", "-")],
+            ylabel="time-cost",
+            plot_name=f"INT8GELU<BLSZ={CONST_BLOCK}><SL={SL}>",
+            args={"SL": SL, "B": CONST_BLOCK, "provider": "triton", "mode": "time-consuming"},
+        )
+    )
+
+
+@triton.testing.perf_report(configs)
+def bench_load_store(
+    SL, CDIM, B, provider, dtype, mode="forward"
+):  # I only use triton as the provider, and mode when benchmarking
+    # create data
+    x = torch.randn(SL, CDIM, dtype=torch.float32).cuda()
+    x = x.to(dtype)
+
+    quantiles = [0.5, 0.2, 0.8]
+    # utility functions
+    if provider == "triton":
+
+        def y_fwd():
+            bench_memory_io_forward(x, B)
+
+    if provider == "torch":
+        torch_gelu = torch.nn.GELU()
+
+        def y_fwd():
+            return torch_gelu(x)
+
+    # forward pass
+    if mode == "time-consuming":
+        convert_func = lambda ms: ms
+        ms, min_ms, max_ms = triton.testing.do_bench(y_fwd, quantiles=quantiles, rep=100)
+    # backward pass
+    if mode == "gbps":
+        convert_func = lambda ms: 2 * x.numel() * x.element_size() / ms * 1e-6
+        ms, min_ms, max_ms = triton.testing.do_bench(y_fwd, quantiles=quantiles, rep=100)
+    return convert_func(ms), convert_func(max_ms), convert_func(min_ms)
+
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    torch.set_printoptions(precision=8, linewidth=1600, sci_mode=False, edgeitems=3)
+    bench_load_store.run(print_data=True)

llava/model/coat/activation/real_quantization/_quantize.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, convert_fp8_to_embit, convert_str_to_fp8, get_configs_io_block
+
+"""Quantize Operator"""
+"""Input uses 1 * 16 group quantization"""
+"""Output uses 1 * 16 group quantization"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_quantize_kernel(
+    output_ptr,
+    output_scale_ptr,  # output
+    input_ptr,  # input
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,
+    fp8_max,  # shape
+    input_stride_0,
+    input_stride_1,  # input stride
+    output_stride_0,
+    output_stride_1,  # output stride
+    s_output_stride_0,
+    s_output_stride_1,  # scale of output stride
+    SCALE_MIN_THRES: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+    input = input.to(tl.float32)
+
+    output = tl.reshape(input, (BLOCK_M, BLOCK_SN, QB))
+
+    # Quantize Scale calculation
+    abs_output = tl.abs(output)
+    max_val = tl.max(abs_output, axis=2) + SCALE_MIN_THRES
+    scale_output = max_val / fp8_max
+    scale_output = tl.reshape(scale_output, (BLOCK_M, BLOCK_SN, 1))
+
+    # Quantize
+    output = tl.fdiv(output, scale_output)
+
+    output = output.to(output_ptr.type.element_ty)
+
+    scale_output = scale_output.to(output_scale_ptr.type.element_ty)
+    scale_output = tl.reshape(scale_output, (BLOCK_M, BLOCK_SN))
+    output = tl.reshape(output, (BLOCK_M, BLOCK_N))
+
+    # debug
+    # gelu_output = input
+    # scale_output = scale_input
+
+    # pointers
+    output_block_ptr = tl.make_block_ptr(
+        base=output_ptr,
+        shape=(M, N),
+        strides=(output_stride_0, output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+    scale_output_ptr = tl.make_block_ptr(
+        base=output_scale_ptr,
+        shape=(M, SN),
+        strides=(s_output_stride_0, s_output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_SN),
+        block_shape=(BLOCK_M, BLOCK_SN),
+        order=(1, 0),
+    )
+
+    tl.store(output_block_ptr, output, boundary_check=(0, 1))
+    tl.store(scale_output_ptr, scale_output, boundary_check=(0, 1))
+
+
+def fp8_quantize(x, QB, fp8type):
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x.shape) == 3:
+        batched = True
+        BS = x.shape[0]
+        x = x.reshape(-1, x.shape[-1])
+
+    # defining the input and output tensor
+    M, N = x.shape
+    SN = N // QB
+
+    if isinstance(fp8type, str):
+        fp8type = convert_str_to_fp8[fp8type]
+    y = torch.empty_like(x, dtype=fp8type)
+    s_y = torch.empty((M, SN), dtype=torch.bfloat16, device=x.device)
+    fp8MaxValue = FP8_MAX_VALUE[fp8type]  # E4M3 and E5M2 have different max value
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_quantize_kernel[grid](
+        y,
+        s_y,
+        x,
+        M,
+        N,
+        SN,
+        QB,
+        fp8MaxValue,
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+        s_y.stride(0),
+        s_y.stride(1),
+        SCALE_MIN_THRES=SCALE_MIN_THRES,
+    )
+
+    # Recover 2D to 3D
+    if batched:
+        y = y.reshape(BS, -1, y.shape[-1])
+        s_y = s_y.reshape(BS, -1, s_y.shape[-1])
+
+    return y, s_y

llava/model/coat/activation/real_quantization/_quantize_pertensor.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from ._division import fp8_division
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, convert_str_to_fp8, get_configs_io_block
+
+"""Per Tensor Quantize Operator"""
+"""Input uses full precision"""
+"""Output uses per tensor quantization"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_quantize_pertensor_kernel(
+    output_scale_ptr,  # output
+    input_ptr,  # input
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,
+    fp8_max,  # shape
+    input_stride_0,
+    input_stride_1,  # input stride
+    s_output_stride_0,
+    s_output_stride_1,  # scale of output stride
+    SCALE_MIN_THRES: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+    input = input.to(tl.float32)
+
+    output = tl.reshape(input, (BLOCK_M, BLOCK_SN, QB))
+
+    # Quantize Scale calculation
+    abs_output = tl.abs(output)
+    max_val = tl.max(abs_output, axis=2) + SCALE_MIN_THRES
+    scale_output = max_val / fp8_max
+    scale_output = tl.reshape(scale_output, (BLOCK_M, BLOCK_SN, 1))
+
+    scale_output = scale_output.to(output_scale_ptr.type.element_ty)
+    scale_output = tl.reshape(scale_output, (BLOCK_M, BLOCK_SN))
+
+    scale_output_ptr = tl.make_block_ptr(
+        base=output_scale_ptr,
+        shape=(M, SN),
+        strides=(s_output_stride_0, s_output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_SN),
+        block_shape=(BLOCK_M, BLOCK_SN),
+        order=(1, 0),
+    )
+
+    tl.store(scale_output_ptr, scale_output, boundary_check=(0, 1))
+
+
+def fp8_quantize_pertensor(x, QB, fp8type, stochastic=False):
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x.shape) == 3:
+        batched = True
+        BS = x.shape[0]
+        x = x.reshape(-1, x.shape[-1])
+
+    # defining the input and output tensor
+    M, N = x.shape
+    SN = N // QB
+
+    fp8type = convert_str_to_fp8[fp8type]
+    s_y = torch.empty((M, SN), dtype=torch.bfloat16, device=x.device)
+    fp8MaxValue = FP8_MAX_VALUE[fp8type]  # E4M3 and E5M2 have different max value
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_quantize_pertensor_kernel[grid](
+        s_y,
+        x,
+        M,
+        N,
+        SN,
+        QB,
+        fp8MaxValue,
+        x.stride(0),
+        x.stride(1),
+        s_y.stride(0),
+        s_y.stride(1),
+        SCALE_MIN_THRES=SCALE_MIN_THRES,
+    )
+
+    s_y_max = s_y.max()
+    y, s_y_max = fp8_division(x, QB, fp8type, s_y_max, stochastic=stochastic)  # reuse the floating point output y1
+
+    # Recover 2D to 3D
+    if batched:
+        y = y.reshape(BS, -1, y.shape[-1])
+        s_y = s_y.reshape(BS, -1, s_y.shape[-1])
+
+    return y, s_y_max, s_y

llava/model/coat/activation/real_quantization/_quantize_pertensor_transpose.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from ._division_transpose import fp8_division_transpose
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, convert_str_to_fp8, get_configs_io_block
+
+"""Per Tensor Quantize and Transpose Operator"""
+"""Input uses floating point tensor"""
+"""Output uses per-tensor quantization, returns a non-transpose version and a transpose version"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_quantize_pertensor_transpose_kernel(
+    output_scale_ptr,  # output
+    input_ptr,  # input
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,
+    fp8_max,  # shape
+    input_stride_0,
+    input_stride_1,  # input stride
+    s_output_stride_0,
+    s_output_stride_1,  # scale of output stride
+    SCALE_MIN_THRES: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+    input = input.to(tl.float32)
+
+    output = tl.reshape(input, (BLOCK_M, BLOCK_SN, QB))
+
+    # Quantize Scale calculation
+    abs_output = tl.abs(output)
+    max_val = tl.max(abs_output, axis=2) + SCALE_MIN_THRES
+    scale_output = max_val / fp8_max
+    scale_output = tl.reshape(scale_output, (BLOCK_M, BLOCK_SN, 1))
+
+    scale_output = scale_output.to(output_scale_ptr.type.element_ty)
+    scale_output = tl.reshape(scale_output, (BLOCK_M, BLOCK_SN))
+
+    scale_output_ptr = tl.make_block_ptr(
+        base=output_scale_ptr,
+        shape=(M, SN),
+        strides=(s_output_stride_0, s_output_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_SN),
+        block_shape=(BLOCK_M, BLOCK_SN),
+        order=(1, 0),
+    )
+
+    tl.store(scale_output_ptr, scale_output, boundary_check=(0, 1))
+
+
+def fp8_quantize_pertensor_transpose(x, QB, fp8type, transpose_output_2d=False, stochastic=False):
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x.shape) == 3:
+        batched = True
+        BS = x.shape[0]
+        x = x.reshape(-1, x.shape[-1])
+
+    # defining the input and output tensor
+    M, N = x.shape
+    SN = N // QB
+
+    fp8type = convert_str_to_fp8[fp8type]
+    s_y = torch.empty((M, SN), dtype=torch.bfloat16, device=x.device)
+    fp8MaxValue = FP8_MAX_VALUE[fp8type]  # E4M3 and E5M2 have different max value
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_quantize_pertensor_transpose_kernel[grid](
+        s_y,
+        x,
+        M,
+        N,
+        SN,
+        QB,
+        fp8MaxValue,
+        x.stride(0),
+        x.stride(1),
+        s_y.stride(0),
+        s_y.stride(1),
+        SCALE_MIN_THRES=SCALE_MIN_THRES,
+    )
+
+    s_y_max = s_y.max()
+    qy, s_y_max, qy_t = fp8_division_transpose(
+        x, QB, fp8type, s_y_max, stochastic=stochastic
+    )  # Stochastic Rounding happens here
+
+    # Recover 2D to 3D
+    if batched:
+        qy = qy.reshape(BS, -1, qy.shape[-1])
+        if not transpose_output_2d:
+            qy_t = qy_t.reshape(BS, -1, qy_t.shape[-1])
+
+    return qy, s_y_max, qy_t  # y_t is expected to be 2D tensor

llava/model/coat/activation/real_quantization/_transpose.py

@@ -0,0 +1,121 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from .common import get_configs_io_block
+
+"""Quantize Operator"""
+"""Input uses 1 * 16 group quantization"""
+"""Output uses 1 * 16 group quantization"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.jit
+def _fp8_transpose_kernel(
+    output_ptr,  # output
+    input_ptr,  # input
+    M,
+    N,  # shape
+    input_stride_0,
+    input_stride_1,  # input stride
+    output_stride_0,
+    output_stride_1,  # output stride
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # pointers
+    input_block_ptr = tl.make_block_ptr(
+        base=input_ptr,
+        shape=(M, N),
+        strides=(input_stride_0, input_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input = tl.load(input_block_ptr)
+
+    output = tl.trans(input)
+
+    # pointers
+    output_block_ptr = tl.make_block_ptr(
+        base=output_ptr,
+        shape=(N, M),
+        strides=(output_stride_0, output_stride_1),
+        offsets=(pid_dim1 * BLOCK_N, pid_dim0 * BLOCK_M),
+        block_shape=(BLOCK_N, BLOCK_M),
+        order=(1, 0),
+    )
+
+    tl.store(output_block_ptr, output, boundary_check=(0, 1))
+
+
+def fp8_transpose(x, transpose_output_2d=False):
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x.shape) == 3:
+        batched = True
+        BS = x.shape[0]
+        x = x.reshape(-1, x.shape[-1])
+
+    # defining the input and output tensor
+    M, N = x.shape
+
+    y = torch.empty((N, M), dtype=x.dtype, device=x.device)
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_transpose_kernel[grid](
+        y,
+        x,
+        M,
+        N,
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+    )
+
+    # Recover 2D to 3D
+    if batched and not transpose_output_2d:
+        y = y.reshape(BS, -1, y.shape[-1])
+
+    return y

llava/model/coat/activation/real_quantization/add_bwd.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from ._division import fp8_division
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, convert_str_to_fp8, get_configs_io_block
+
+"""Element-wise Add, useful for backward"""
+"""Input1 (Residual) uses full-precision/BF16"""
+"""Input2 (Backbone) uses full-precision/BF16"""
+"""Output1 uses full-precision/BF16"""
+"""Output2 uses per-tensor quantization"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_add_Ifp_Ifp_Ofp_Opt_kernel(
+    output1_ptr,  # output
+    output2_scale_ptr,
+    input1_ptr,  # input
+    input2_ptr,  # input
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,
+    fp8_max,  # shape
+    input1_stride_0,
+    input1_stride_1,  # input1 stride
+    input2_stride_0,
+    input2_stride_1,  # input2 stride
+    output1_stride_0,
+    output1_stride_1,  # output stride
+    s_output2_stride_0,
+    s_output2_stride_1,  # scale of output stride
+    SCALE_MIN_THRES: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # --- The first input ---
+    input1_block_ptr = tl.make_block_ptr(
+        base=input1_ptr,
+        shape=(M, N),
+        strides=(input1_stride_0, input1_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input1 = tl.load(input1_block_ptr)
+    input1 = input1.to(tl.float32)
+    input1 = tl.reshape(input1, (BLOCK_M, BLOCK_SN, QB))
+
+    # --- The second input ---
+    input2_block_ptr = tl.make_block_ptr(
+        base=input2_ptr,
+        shape=(M, N),
+        strides=(input2_stride_0, input2_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input2 = tl.load(input2_block_ptr)
+    input2 = input2.to(tl.float32)
+    input2 = tl.reshape(input2, (BLOCK_M, BLOCK_SN, QB))
+
+    # Actual Calculation of Add
+    add_output = input1 + input2
+
+    # Quantize the grad 1 - Scale calculation
+    abs_add_output = tl.abs(add_output)
+    max_val = tl.max(abs_add_output, axis=2) + SCALE_MIN_THRES
+    scale_output2 = max_val / fp8_max
+    scale_output2 = tl.reshape(scale_output2, (BLOCK_M, BLOCK_SN, 1))
+
+    # save the fp add output
+    fp_add_output = add_output.to(output1_ptr.type.element_ty)
+    fp_add_output = tl.reshape(fp_add_output, (BLOCK_M, BLOCK_N))
+
+    # pointers
+    output1_block_ptr = tl.make_block_ptr(
+        base=output1_ptr,
+        shape=(M, N),
+        strides=(output1_stride_0, output1_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    tl.store(output1_block_ptr, fp_add_output, boundary_check=(0, 1))
+
+    # Quantize
+    scale_output2 = scale_output2.to(output2_scale_ptr.type.element_ty)
+    scale_output2 = tl.reshape(scale_output2, (BLOCK_M, BLOCK_SN))
+
+    # pointers
+    scale_output2_ptr = tl.make_block_ptr(
+        base=output2_scale_ptr,
+        shape=(M, SN),
+        strides=(s_output2_stride_0, s_output2_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_SN),
+        block_shape=(BLOCK_M, BLOCK_SN),
+        order=(1, 0),
+    )
+    tl.store(scale_output2_ptr, scale_output2, boundary_check=(0, 1))
+
+
+def fp8_add_Ifp_Ifp_Ofp_Opt(x1, x2, QB, fp8type, stochastic=False):  # suppose x1 is full precision or BF16
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x1.shape) == 3:
+        assert len(x2.shape) == 3
+        batched = True
+        BS = x1.shape[0]
+        x1 = x1.reshape(-1, x1.shape[-1])
+        x2 = x2.reshape(-1, x2.shape[-1])
+
+    # defining the input and output tensor
+    M, N = x1.shape
+    SN = N // QB
+    assert x1.shape == x2.shape
+
+    if isinstance(fp8type, str):
+        fp8type = convert_str_to_fp8[fp8type]
+    y1 = torch.empty_like(x1, dtype=torch.bfloat16)
+    s_y2 = torch.empty((M, SN), dtype=torch.bfloat16, device=x2.device)
+    fp8MaxValue = FP8_MAX_VALUE[fp8type]  # E4M3 and E5M2 have different max value
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_add_Ifp_Ifp_Ofp_Opt_kernel[grid](
+        y1,
+        s_y2,
+        x1,
+        x2,
+        M,
+        N,
+        SN,
+        QB,
+        fp8MaxValue,
+        x1.stride(0),
+        x1.stride(1),
+        x2.stride(0),
+        x2.stride(1),
+        y1.stride(0),
+        y1.stride(1),
+        s_y2.stride(0),
+        s_y2.stride(1),
+        SCALE_MIN_THRES=SCALE_MIN_THRES,
+    )
+
+    s_y2_max = s_y2.max()
+    qy2, s_y2_max = fp8_division(y1, QB, fp8type, s_y2_max, stochastic=stochastic)  # reuse the floating point output y1
+
+    # Recover 2D to 3D
+    if batched:
+        y1 = y1.reshape(BS, -1, y1.shape[-1])
+        qy2 = qy2.reshape(BS, -1, qy2.shape[-1])
+        s_y2 = s_y2.reshape(BS, -1, s_y2.shape[-1])
+
+    return y1, (qy2, s_y2_max, s_y2)

llava/model/coat/activation/real_quantization/add_fwd.py

@@ -0,0 +1,219 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
+# LICENSE is in incl_licenses directory.
+
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+# 4 block
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+from .common import FP8_MAX_VALUE, SCALE_MIN_THRES, get_configs_io_block
+
+"""Element-wise Add, used in forward pass"""
+"""Input1 (Residual) uses full-precision/BF16"""
+"""Input2 (Backbone) uses full-precision/BF16"""
+"""Output1 uses full-precision/BF16"""
+"""Output2 uses 1 * 16 group quantization"""
+"""The input can be 2D or 3D, but the calculation is performed in 2D"""
+
+
+@triton.autotune(
+    configs=[] + get_configs_io_block(),
+    key=[
+        "N",
+    ],
+)
+@triton.heuristics(
+    {
+        "BLOCK_SN": lambda args: args["BLOCK_N"] // args["QB"],
+    }
+)
+@triton.jit
+def _fp8_add_Ifp_Ifp_Ofp_Og16_kernel(
+    output1_ptr,  # output
+    output2_ptr,
+    output2_scale_ptr,
+    input1_ptr,  # input
+    input2_ptr,  # input
+    M,
+    N,
+    SN,
+    QB: tl.constexpr,
+    fp8_max,  # shape
+    input1_stride_0,
+    input1_stride_1,  # input1 stride
+    input2_stride_0,
+    input2_stride_1,  # input2 stride
+    output1_stride_0,
+    output1_stride_1,  # output stride
+    output2_stride_0,
+    output2_stride_1,  # output stride
+    s_output2_stride_0,
+    s_output2_stride_1,  # scale of output stride
+    SCALE_MIN_THRES: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_SN: tl.constexpr,
+):  # CUDA block size
+
+    # Block PID
+    pid = tl.program_id(0)
+    NUM_BLOCK_N = tl.cdiv(N, BLOCK_N)
+    pid_dim0 = pid // NUM_BLOCK_N
+    pid_dim1 = pid % NUM_BLOCK_N
+
+    # --- The first input ---
+    input1_block_ptr = tl.make_block_ptr(
+        base=input1_ptr,
+        shape=(M, N),
+        strides=(input1_stride_0, input1_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input1 = tl.load(input1_block_ptr)
+    input1 = input1.to(tl.float32)
+    input1 = tl.reshape(input1, (BLOCK_M, BLOCK_SN, QB))
+
+    # --- The second input ---
+    input2_block_ptr = tl.make_block_ptr(
+        base=input2_ptr,
+        shape=(M, N),
+        strides=(input2_stride_0, input2_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    input2 = tl.load(input2_block_ptr)
+    input2 = input2.to(tl.float32)
+    input2 = tl.reshape(input2, (BLOCK_M, BLOCK_SN, QB))
+
+    # Actual Calculation of Add
+    add_output = input1 + input2
+
+    # Quantize the grad 1 - Scale calculation
+    abs_add_output = tl.abs(add_output)
+    max_val = tl.max(abs_add_output, axis=2) + SCALE_MIN_THRES
+    scale_output2 = max_val / fp8_max
+    scale_output2 = tl.reshape(scale_output2, (BLOCK_M, BLOCK_SN, 1))
+
+    # save the fp add output
+    fp_add_output = add_output.to(output1_ptr.type.element_ty)
+    fp_add_output = tl.reshape(fp_add_output, (BLOCK_M, BLOCK_N))
+
+    # pointers
+    output1_block_ptr = tl.make_block_ptr(
+        base=output1_ptr,
+        shape=(M, N),
+        strides=(output1_stride_0, output1_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    tl.store(output1_block_ptr, fp_add_output)
+
+    # Quantize
+    add_output = tl.fdiv(add_output, scale_output2)
+    scale_output2 = scale_output2.to(output2_scale_ptr.type.element_ty)
+    scale_output2 = tl.reshape(scale_output2, (BLOCK_M, BLOCK_SN))
+    add_output = tl.reshape(add_output, (BLOCK_M, BLOCK_N))
+
+    add_output = add_output.to(output2_ptr.type.element_ty)
+    add_output = tl.reshape(add_output, (BLOCK_M, BLOCK_N))
+
+    # pointers
+    output2_block_ptr = tl.make_block_ptr(
+        base=output2_ptr,
+        shape=(M, N),
+        strides=(output2_stride_0, output2_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_N),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+    scale_output2_ptr = tl.make_block_ptr(
+        base=output2_scale_ptr,
+        shape=(M, SN),
+        strides=(s_output2_stride_0, s_output2_stride_1),
+        offsets=(pid_dim0 * BLOCK_M, pid_dim1 * BLOCK_SN),
+        block_shape=(BLOCK_M, BLOCK_SN),
+        order=(1, 0),
+    )
+    tl.store(output2_block_ptr, add_output, boundary_check=(0, 1))
+    tl.store(scale_output2_ptr, scale_output2, boundary_check=(0, 1))
+
+
+def fp8_add_Ifp_Ifp_Ofp_Og16(x1, x2, fp8type, QB):  # suppose x1 is full precision or BF16
+    # Change batched 3D input to 2D
+    batched = False
+    if len(x1.shape) == 3:
+        batched = True
+        BS = x1.shape[0]
+        x1 = x1.reshape(-1, x1.shape[-1])
+        x2 = x2.reshape(-1, x2.shape[-1])
+
+    # defining the input and output tensor
+    M, N = x1.shape
+    SN = int(N / QB)  # assume the shape of quantization block size is always 1 * G
+    assert x1.shape == x2.shape
+
+    y1 = torch.empty_like(x1, dtype=torch.bfloat16)
+    y2 = torch.empty_like(x2, dtype=fp8type)
+    s_y2 = torch.empty((M, SN), dtype=torch.bfloat16, device=x2.device)
+    fp8MaxValue = FP8_MAX_VALUE[fp8type]  # E4M3 and E5M2 have different max value
+
+    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
+
+    _fp8_add_Ifp_Ifp_Ofp_Og16_kernel[grid](
+        y1,
+        y2,
+        s_y2,
+        x1,
+        x2,
+        M,
+        N,
+        SN,
+        QB,
+        fp8MaxValue,
+        x1.stride(0),
+        x1.stride(1),
+        x2.stride(0),
+        x2.stride(1),
+        y1.stride(0),
+        y1.stride(1),
+        y2.stride(0),
+        y2.stride(1),
+        s_y2.stride(0),
+        s_y2.stride(1),
+        SCALE_MIN_THRES=SCALE_MIN_THRES,
+    )
+
+    # Recover 2D to 3D
+    if batched:
+        y1 = y1.reshape(BS, -1, y1.shape[-1])
+        y2 = y2.reshape(BS, -1, y2.shape[-1])
+        s_y2 = s_y2.reshape(BS, -1, s_y2.shape[-1])
+
+    return y1, (y2, s_y2)