Upload files with `vila-upload`.

Upload __init__.py
Upload utils.py
Upload config.json
Upload README.md
Upload model_utils_packing.py
Upload auto_processor.py
Upload modeling_vila.py
Upload loss.py
Upload distributed.py
Upload llm/model.safetensors
Upload llm/config.json
Upload mm_projector/model.safetensors
Upload mm_projector/config.json
Upload vision_tower/model.safetensors
Upload vision_tower/config.json

README.md

@@ -52,34 +52,36 @@ print(colored(response, "cyan", attrs=["bold"]))
 
 we also support `AutoProcessor` class to ease data preparation for training and finetuning.
 
+
+### single call
+
 ```python
 from transformers import AutoProcessor, AutoModel
 
 model_path = "Efficient-Large-Model/NVILA-Lite-2B-hf-preview"
 processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
+# important: set model to eval mode, otherwise the model will be in training mode and will pad to right.
+model.eval()
 
-gpt_conv = [ {
+gpt_conv = [{
     "role": "user",
     "content": [
         {"type": "image", "path": "demo_images/demo_img_1.png"},
         {"type": "text", "text": "Describe this image."}
     ]
 }]
+text = processor.apply_chat_template(gpt_conv, tokenize=False, add_generation_prompt=True)
+inputs = processor([text])
 
-inputs = processor.apply_chat_template(conversation=gpt_conv, padding=True, return_tensors="pt")
-model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
 output_ids = model.generate(
     input_ids=inputs.input_ids,
-    media={
-        "image": inputs.image,
-    },
-    media_config={
-        "image": {}
-    },
+    media=inputs.media,
+    media_config=inputs.media_config,
     generation_config=model.generation_config,
     max_new_tokens=256,
 )
-print(processor.tokenizer.decode(output_ids[0], skip_special_tokens=True))
+print(processor.tokenizer.batch_decode(output_ids, skip_special_tokens=True))
 
 ##### the above code is equivalent to
 # response = model.generate_content([
@@ -89,6 +91,54 @@ print(processor.tokenizer.decode(output_ids[0], skip_special_tokens=True))
 # print(colored(response, "cyan", attrs=["bold"]))
 ```
 
+### batch call
+
+```python
+from transformers import AutoProcessor, AutoModel
+
+model_path = "Efficient-Large-Model/NVILA-Lite-2B-hf-preview"
+model_path = "./NVILA-Lite-2B-hf-preview"
+processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
+# important: set model to eval mode, otherwise the model will be in training mode and will pad to right.
+model.eval()
+
+gpt_conv1 = [{
+    "role": "user",
+    "content": [
+        {"type": "image", "path": "demo_images/demo_img_1.png"},
+        {"type": "text", "text": "Describe this image."}
+    ]
+}]
+gpt_conv2 = [{
+    "role": "user",
+    "content": [
+        {"type": "image", "path": "demo_images/demo_img_2.png"},
+        {"type": "text", "text": "Describe this image for me. Provide a detailed description of the image."}
+    ]
+}]
+
+messages = [gpt_conv1, gpt_conv2]
+texts = [
+    processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+    for msg in messages
+]
+inputs = processor(texts)
+
+output_ids = model.generate(
+    input_ids=inputs.input_ids,
+    media=inputs.media,
+    media_config=inputs.media_config,
+    generation_config=model.generation_config,
+    max_new_tokens=256,
+)
+output_texts = processor.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+print(output_texts[0])
+print("---" * 40)
+print(output_texts[1])
+```
+
+
 ## Model Convert
 
 The follwing code converts a convetional NVILA model to a HF compatible model.

auto_processor.py

@@ -1,8 +1,11 @@
+import copy
 import os
 import os.path as osp
+import warnings
 from collections import defaultdict
 from typing import List, Union
 
+import torch
 from transformers import AutoConfig, AutoImageProcessor, AutoModel, AutoProcessor, AutoTokenizer
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.image_utils import ImageInput, VideoInput
@@ -16,6 +19,26 @@ from .mm_utils import process_image, process_images
 from .tokenizer_utils import tokenize_conversation
 
 
+def vila_pad_fn(input_ids_list, padding_value=0, target_len=None, padding_side="left"):
+    # tensor shape is (batch_size, seq_len)
+    max_len = max([ids.shape[1] for ids in input_ids_list])
+    if target_len is not None:
+        assert target_len >= max_len, "target_len must be greater than or equal to max_len"
+        max_len = target_len
+
+    new_input_ids_list = []
+    for i, input_ids in enumerate(input_ids_list):
+        pad_tensor = torch.ones_like(input_ids) * padding_value
+        curr_len = input_ids.shape[1]
+        pad_tensor = pad_tensor[:, : max_len - curr_len]
+        if padding_side == "right":
+            input_ids = torch.cat((input_ids, pad_tensor), dim=1)
+        else:
+            input_ids = torch.cat((pad_tensor, input_ids), dim=1)
+        new_input_ids_list.append(input_ids)
+    return torch.cat(new_input_ids_list, dim=0)
+
+
 class VILAProcessorKwargs(ProcessingKwargs, total=False):
     _defaults = {
         "text_kwargs": {
@@ -40,6 +63,7 @@ class VILAProcessor(ProcessorMixin):
         self.config = config
         self.image_processor = image_processor
         self.tokenizer = tokenizer
+
         super().__init__(image_processor, tokenizer, chat_template=chat_template)
 
     @classmethod
@@ -48,7 +72,7 @@ class VILAProcessor(ProcessorMixin):
             pretrained_model_name_or_path = pretrained_model_name_or_path
         else:
             print(f"pretrained_model_name_or_path {pretrained_model_name_or_path} is not a directory, downloading")
-            from huggingface_hub import HfApi, snapshot_download
+            from huggingface_hub import snapshot_download
 
             pretrained_model_name_or_path = snapshot_download(pretrained_model_name_or_path)
 
@@ -59,7 +83,6 @@ class VILAProcessor(ProcessorMixin):
             osp.join(pretrained_model_name_or_path, "llm"), trust_remote_code=True
         )
         config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
-
         return cls(image_processor=image_processor, tokenizer=tokenizer, config=config)
 
     def __repr__(self):
@@ -74,9 +97,45 @@ class VILAProcessor(ProcessorMixin):
         text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
         videos: VideoInput = None,
         **kwargs: Unpack[VILAProcessorKwargs],
+    ) -> BatchFeature:
+        if images is not None:
+            warnings.warn("images is not supported in __call__")
+
+        input_ids = []
+        media = defaultdict(list)
+        media_config = defaultdict(dict)
+        for conv in conversation:
+            feat = self.__single_call__(conv, images, text, videos, **kwargs)
+            input_ids.append(feat.input_ids)
+            for name in feat.media:
+                media[name] += feat.media[name]
+            for name in feat.media_config:
+                media_config[name].update(feat.media_config[name])
+
+        return BatchFeature(
+            data={
+                # "input_ids": torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=self.pad_token_id),
+                "input_ids": vila_pad_fn(
+                    input_ids,
+                    padding_value=self.tokenizer.pad_token_id,
+                    padding_side="left",
+                ),
+                "media": media,
+                "media_config": media_config,
+            }
+        )
+
+    def __single_call__(
+        self,
+        conversation,
+        images: ImageInput = None,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+        videos: VideoInput = None,
+        **kwargs: Unpack[VILAProcessorKwargs],
     ) -> BatchFeature:
         # TODO: should be merged with llava_arch.py/generate_content()
         # TODO (extract and preprocess should be done together, as the preprocess of image and video can be different, i.e. when dynamic res is used)
+        conversation = copy.deepcopy(conversation)
         media = extract_media(conversation, self.config)
         # Process media
         media_config = defaultdict(dict)
@@ -107,11 +166,9 @@ class VILAProcessor(ProcessorMixin):
                 ]
             else:
                 raise ValueError(f"Unsupported media type: {name}")
-
         input_ids = tokenize_conversation(conversation, self.tokenizer, add_generation_prompt=True).cuda().unsqueeze(0)
         # Set up the generation config
-        # print(input_ids.shape); print(media); input()
-        return BatchFeature(data={"input_ids": input_ids, **media})
+        return BatchFeature(data={"input_ids": input_ids, "media": media, "media_config": media_config})
 
     def batch_decode(self, *args, **kwargs):
         """
@@ -152,7 +209,6 @@ class VILAProcessor(ProcessorMixin):
     #     inputs = processor(conversation=llavaconv, padding=True, return_tensors="pt")
     def apply_chat_template(self, conversation, add_generation_prompt=True, **kwargs):
         vila_conv = []
-
         for chat in conversation:
             vila_chat = {"from": "", "value": []}
             if chat["role"] == "user":
@@ -172,7 +228,8 @@ class VILAProcessor(ProcessorMixin):
                     vila_chat["value"].append(content["text"])
             vila_conv.append(vila_chat)
 
-        return self(vila_conv)
+        # return self(vila_conv)
+        return vila_conv
 
 
 if __name__ == "__main__":

config.json

@@ -1,19 +1,26 @@
 {
   "_attn_implementation_autoset": true,
-  "_name_or_path": "runs/train/qwen25_2B_3x3-sft-20241118122815/model",
+  "_name_or_path": "NVILA-Lite-2B-hf-preview",
   "architectures": [
     "VILAForCasualLM"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_vila.VILAConfig",
+    "AutoModel": "modeling_vila.VILAForCasualLM",
+    "AutoModelForCausalLM": "modeling_vila.VILAForCasualLM",
+    "AutoProcessor": "auto_processor.VILAProcessor"
+  },
   "chat_template": null,
   "drop_path_rate": 0.0,
   "dynamic_s2": false,
   "fps": 0.0,
   "hidden_size": 1536,
   "image_aspect_ratio": "dynamic",
+  "image_encoder": "{\"_target_\": \"llava.model.encoders.BasicImageEncoder\"}",
   "interpolate_mode": "linear",
   "llm_cfg": {
     "_attn_implementation_autoset": false,
-    "_name_or_path": "runs/train/qwen25_2B_3x3-sft-20241118122815/model/llm",
+    "_name_or_path": "NVILA-Lite-2B-hf-preview/llm",
     "add_cross_attention": false,
     "architectures": [
       "Qwen2ForCausalLM"
@@ -53,7 +60,7 @@
     "max_position_embeddings": 32768,
     "max_window_layers": 28,
     "min_length": 0,
-    "model_max_length": 4096,
+    "model_max_length": null,
     "model_type": "qwen2",
     "no_repeat_ngram_size": 0,
     "num_attention_heads": 12,
@@ -89,18 +96,20 @@
     "tokenizer_padding_side": "right",
     "top_k": 50,
     "top_p": 1.0,
-    "torch_dtype": "bfloat16",
+    "torch_dtype": "float16",
     "torchscript": false,
     "typical_p": 1.0,
     "use_bfloat16": false,
     "use_cache": true,
     "use_sliding_window": false,
-    "vocab_size": 151648
+    "vocab_size": 151651
   },
+  "max_tiles": 12,
+  "min_tiles": 1,
   "mm_hidden_size": 1152,
   "mm_projector_cfg": {
     "_attn_implementation_autoset": false,
-    "_name_or_path": "runs/train/qwen25_2B_3x3-sft-20241118122815/model/mm_projector",
+    "_name_or_path": "NVILA-Lite-2B-hf-preview/mm_projector",
     "add_cross_attention": false,
     "architectures": [
       "MultimodalProjector"
@@ -160,7 +169,7 @@
     "tokenizer_class": null,
     "top_k": 50,
     "top_p": 1.0,
-    "torch_dtype": "bfloat16",
+    "torch_dtype": "float16",
     "torchscript": false,
     "typical_p": 1.0,
     "use_bfloat16": false
@@ -186,10 +195,12 @@
   "tune_language_model": true,
   "tune_mm_projector": true,
   "tune_vision_tower": true,
+  "version": "2.0",
+  "video_encoder": "{\"_target_\": \"llava.model.encoders.BasicVideoEncoder\"}",
   "vision_resolution": -1,
   "vision_tower_cfg": {
     "_attn_implementation_autoset": false,
-    "_name_or_path": "runs/train/qwen25_2B_3x3-sft-20241118122815/model/vision_tower",
+    "_name_or_path": "NVILA-Lite-2B-hf-preview/vision_tower",
     "add_cross_attention": false,
     "architectures": [
       "SiglipVisionModel"
@@ -261,17 +272,11 @@
     "tokenizer_class": null,
     "top_k": 50,
     "top_p": 1.0,
-    "torch_dtype": "bfloat16",
+    "torch_dtype": "float16",
     "torchscript": false,
     "typical_p": 1.0,
     "use_bfloat16": false,
     "vision_use_head": false
   },
-  "version": "2.0",
-  "auto_map": {
-    "AutoProcessor": "auto_processor.VILAProcessor",
-    "AutoConfig": "modeling_vila.VILAConfig",
-    "AutoModel": "modeling_vila.VILAForCasualLM",
-    "AutoModelForCausalLM": "modeling_vila.VILAForCasualLM"
-  }
-}
+  "vision_tower_lr": null
+}

distributed.py

@@ -0,0 +1,73 @@
+import os
+import warnings
+from typing import Any, List, Optional
+
+from torch import distributed as dist
+
+__all__ = [
+    "init",
+    "is_initialized",
+    "size",
+    "rank",
+    "local_size",
+    "local_rank",
+    "is_main",
+    "barrier",
+    "gather",
+    "all_gather",
+]
+
+
+def init() -> None:
+    if "RANK" not in os.environ:
+        warnings.warn("Environment variable `RANK` is not set. Skipping distributed initialization.")
+        return
+    dist.init_process_group(backend="nccl", init_method="env://")
+
+
+def is_initialized() -> bool:
+    return dist.is_initialized()
+
+
+def size() -> int:
+    return int(os.environ.get("WORLD_SIZE", 1))
+
+
+def rank() -> int:
+    return int(os.environ.get("RANK", 0))
+
+
+def local_size() -> int:
+    return int(os.environ.get("LOCAL_WORLD_SIZE", 1))
+
+
+def local_rank() -> int:
+    return int(os.environ.get("LOCAL_RANK", 0))
+
+
+def is_main() -> bool:
+    return rank() == 0
+
+
+def barrier() -> None:
+    dist.barrier()
+
+
+def gather(obj: Any, dst: int = 0) -> Optional[List[Any]]:
+    if not is_initialized():
+        return [obj]
+    if is_main():
+        objs = [None for _ in range(size())]
+        dist.gather_object(obj, objs, dst=dst)
+        return objs
+    else:
+        dist.gather_object(obj, dst=dst)
+        return None
+
+
+def all_gather(obj: Any) -> List[Any]:
+    if not is_initialized():
+        return [obj]
+    objs = [None for _ in range(size())]
+    dist.all_gather_object(objs, obj)
+    return objs

llm/config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "runs/train/qwen25_2B_3x3-sft-20241118122815/model/llm",
+  "_name_or_path": "NVILA-Lite-2B-hf-preview/llm",
   "architectures": [
     "Qwen2ForCausalLM"
   ],
@@ -12,7 +12,7 @@
   "intermediate_size": 8960,
   "max_position_embeddings": 32768,
   "max_window_layers": 28,
-  "model_max_length": 4096,
+  "model_max_length": null,
   "model_type": "qwen2",
   "num_attention_heads": 12,
   "num_hidden_layers": 28,
@@ -24,9 +24,9 @@
   "tie_word_embeddings": true,
   "tokenizer_model_max_length": 4096,
   "tokenizer_padding_side": "right",
-  "torch_dtype": "bfloat16",
+  "torch_dtype": "float16",
   "transformers_version": "4.46.0",
   "use_cache": true,
   "use_sliding_window": false,
-  "vocab_size": 151648
+  "vocab_size": 151651
 }

llm/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fe90977102f7c2ffc3c0c9fff9cb6bad16937a2e93c49a0a41976fc2a50dd077
-size 3086582408
+oid sha256:e025b53d85054b09b0dc887b14a07d8e1faeda7430842fbf6940f1c1177adf3a
+size 3086591288

loss.py

@@ -0,0 +1,48 @@
+from typing import List, Union
+
+import torch
+from torch.nn.functional import cross_entropy
+
+from .constants import IGNORE_INDEX
+
+__all__ = ["soft_cross_entropy"]
+
+
+def soft_cross_entropy(
+    outputs: torch.Tensor,
+    targets: torch.Tensor,
+    soft_tokens: Union[torch.Tensor, List[int]],
+    std: float = 1,
+    ignore_index: int = IGNORE_INDEX,
+) -> torch.Tensor:
+    # Remove last token from outputs and first token from targets
+    outputs = outputs[..., :-1, :].contiguous()
+    targets = targets[..., 1:].contiguous()
+
+    # Flatten outputs and targets
+    targets = targets.view(-1)
+    outputs = outputs.view(targets.size(0), -1)
+
+    # Remove outputs and targets with ignore_index
+    indices = targets != ignore_index
+    outputs = outputs[indices]
+    targets = targets[indices]
+
+    # Convert soft token IDs to tensor
+    if isinstance(soft_tokens, list):
+        soft_tokens = torch.tensor(soft_tokens).to(targets)
+
+    # Calculate loss for non-soft tokens
+    indices = torch.isin(targets, soft_tokens, invert=True)
+    loss = cross_entropy(outputs[indices], targets[indices], reduction="sum")
+
+    # Calculate loss for soft tokens
+    indices = torch.isin(targets, soft_tokens)
+    targets_indices = torch.zeros_like(outputs[indices])
+    for k, target in enumerate(targets[indices]):
+        dist = torch.exp(-((target - soft_tokens) ** 2) / (2 * std**2))
+        targets_indices[k][soft_tokens] = dist / dist.sum()
+    loss += cross_entropy(outputs[indices], targets_indices, reduction="sum")
+
+    # Return average loss
+    return loss / targets.size(0)

mm_projector/config.json

@@ -1,10 +1,10 @@
 {
-  "_name_or_path": "runs/train/qwen25_2B_3x3-sft-20241118122815/model/mm_projector",
+  "_name_or_path": "NVILA-Lite-2B-hf-preview/mm_projector",
   "architectures": [
     "MultimodalProjector"
   ],
   "mm_projector_type": "mlp_downsample_3x3_fix",
   "model_type": "v2l_projector",
-  "torch_dtype": "bfloat16",
+  "torch_dtype": "float16",
   "transformers_version": "4.46.0"
 }

mm_projector/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:191ffde694a269d7b0ca2f1e30da5d5fecf2a9bb8f4879fbf0b780368c6a9cc4
-size 87068272
+oid sha256:9a884468cfc9b3ff37715e3bdb1bbea378fba5eb7a1b883e2958a73313eb5d35
+size 87068256

model_utils_packing.py

@@ -0,0 +1,35 @@
+from importlib import import_module
+from typing import Tuple
+
+import torch
+import transformers
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ["patch"]
+
+
+def _get_unpad_data(attention_mask: torch.Tensor, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, int]:
+    if hasattr(_get_unpad_data, "seqlens_in_batch"):
+        seqlens_in_batch = _get_unpad_data.seqlens_in_batch
+    else:
+        seqlens_in_batch = torch.sum(attention_mask, dim=1)
+
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return indices, cu_seqlens, max_seqlen_in_batch
+
+
+def set_seqlens_in_batch(seqlens_in_batch: torch.Tensor) -> None:
+    _get_unpad_data.seqlens_in_batch = seqlens_in_batch
+
+
+def patch(model: nn.Module) -> None:
+    if transformers.__version__ < "4.43.0":
+        m = import_module(model.__module__)
+        if not hasattr(m, "_get_unpad_data"):
+            raise ValueError(f"Module {m} does not have function '_get_unpad_data' for packing")
+        m._get_unpad_data = _get_unpad_data
+    else:
+        transformers.modeling_flash_attention_utils._get_unpad_data = _get_unpad_data

modeling_vila.py

@@ -44,15 +44,21 @@ from .builder import build_llm_and_tokenizer
 from .configuration_vila import VILAConfig
 from .constants import *
 from .conversation import SeparatorStyle, default_conversation
+from .distributed import all_gather as vila_all_gather
+from .loss import soft_cross_entropy
 from .media import extract_media
 from .media_encoder import BasicImageEncoder, BasicVideoEncoder
 from .mm_utils import process_image, process_images
+from .model_utils_packing import set_seqlens_in_batch
 from .siglip_encoder import SiglipVisionTower, SiglipVisionTowerDynamicS2, SiglipVisionTowerS2
 from .tokenizer_utils import tokenize_conversation
 from .utils import get_model_config, load_tokenizer_then_handle_media_tokens_and_chat_template
 
-
 # from llava.constants import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX, NUM_EXTRA_TOKENS
+
+# ease debugging
+python_input = input
+
 # quick hack for remote code
 def get_pg_manager():
     return None
@@ -97,28 +103,29 @@ def get_vila_version(model_path: str) -> str:
 
 def generate_jinja_template(conv_mode: str) -> str:
     if conv_mode == "vicuna_v1":
-        return """{% set system_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions." %}
-{% set roles = ["USER", "ASSISTANT"] %}
+        return """{% set system_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. " %}
+{% set roles = ["user", "assistant"] %}
 {% set sep = " " %}
-{% set sep2 = "</s>" %}
 
 {{ system_prompt }}
 
 {% for message in messages %}
     {% if message['role'] == roles[0] %}
-        {{ roles[0] }}{{ sep }}{{ message['content'] }}{{ sep2 }}
+        {{ "USER: " }}{{ sep }}{{ message['content'] }}{{ sep }}
     {% else %}
-        {{ roles[1] }}{{ sep }}{{ message['content'] }}{{ sep2 }}
+        {{ "ASSISTANT: " }}{{ sep }}{{ message['content'] }}{{ sep }}
     {% endif %}
-{% endfor %}"""
+{% endfor %}
+{% if messages[-1]['role'] == 'user' %}
+    {{ "ASSISTANT:" }}
+{% endif %}
+"""
     elif conv_mode == "llama_3":
-        return """{% set system_prompt = "<|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." %}
-{% set roles = ["<|start_header_id|>user<|end_header_id|>\n\n", "<|start_header_id|>assistant<|end_header_id|>\n\n"] %}
+        return """{% set system_prompt = "<|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.<|eot_id|>" %}
+{% set roles = ["<|start_header_id|>user<|end_header_id|>\\n\\n", "<|start_header_id|>assistant<|end_header_id|>\\n\\n"]%}
 {% set sep = "<|eot_id|>" %}
-{% set sep2 = "<|end_of_text|>" %}
 
 {{ system_prompt }}
-
 {% for message in messages %}
     {% if message['role'] == 'user' %}
         {{ roles[0] }}{{ message['content'] }}{{ sep }}
@@ -126,8 +133,10 @@ def generate_jinja_template(conv_mode: str) -> str:
         {{ roles[1] }}{{ message['content'] }}{{ sep }}
     {% endif %}
 {% endfor %}
-
-{{ sep2 }}"""
+{% if messages[-1]['role'] == 'user' %}
+    {{ roles[1] }}
+{% endif %}
+"""
     elif conv_mode == "hermes_2":
         return """{% set system_prompt = "<|im_start|>system\nAnswer the questions." %}
 {% set roles = ["<|im_start|>user\n", "<|im_start|>assistant\n"] %}
@@ -202,6 +211,7 @@ class VILAPretrainedModel(PreTrainedModel):
         # set device_map auto can autoamtically shard llm to different devices
         self.llm, self.tokenizer = self.init_llm(llm_cfg, config, device_map=device_map)
 
+        # NOTE(ligeng): need to add other decoders from config
         self.encoders = {"image": BasicImageEncoder(self), "video": BasicVideoEncoder(self)}
 
         self.post_config()
@@ -218,18 +228,20 @@ class VILAPretrainedModel(PreTrainedModel):
         output_dir: str = None,
         vila_version: str | None = None,
         conv_mode: str | None = None,
-        copy: bool = True,
+        copy: bool = False,
+        copy_weights: bool = True,
+        copy_code: bool = True,
         *model_args,
         **kwargs,
     ):
         # assert type(self) == VILAForCasualLM, "This method is only available for VILAForCasualLM."
-        from huggingface_hub import HfApi, snapshot_download
-
+        assert model_path != output_dir, "model_path and output_dir cannot be the same"
         if os.path.isdir(model_path):
             model_path = model_path
         else:
-            api = HfApi()
-            model_path = snapshot_download(model_path, local_dir=output_dir)
+            from huggingface_hub import HfApi, snapshot_download
+
+            model_path = snapshot_download(model_path)
             print("downloading HF model to", model_path)
 
         if check_dot_in_model_path(model_path) and output_dir is None:
@@ -240,10 +252,18 @@ class VILAPretrainedModel(PreTrainedModel):
             raise ValueError(
                 f"Output directory {output_dir} contains a dot, which will affect the remote code loading. Please specify a valid output directory without dots."
             )
+
+        if copy:
+            print("copy is set to True, copying weights and code to output_dir")
+            copy_weights = copy_code = True
+        # copy weights and code to output_dir
+        self.copy_or_symlink_directory(model_path, output_dir, copy=copy_weights)
+        self.copy_remote_py_files(output_dir, copy=copy_code)
+
         if vila_version is None:
-            vila_version = get_vila_version(model_path)
+            vila_version = get_vila_version(output_dir)
 
-        cfg_path = os.path.join(model_path, "config.json")
+        cfg_path = os.path.join(output_dir, "config.json")
         config = json.load(open(cfg_path))
         config["version"] = "2.0"  # nvila tag
         config["architectures"] = ["VILAForCasualLM"]
@@ -253,24 +273,53 @@ class VILAPretrainedModel(PreTrainedModel):
             "AutoModel": "modeling_vila.VILAForCasualLM",
             "AutoModelForCausalLM": "modeling_vila.VILAForCasualLM",
         }
+        # vila1.5 legacy support
         config["model_type"] = "vila"
         if vila_version in ["vila1.5", "vila-m3"]:
             if conv_mode is None:
-                raise ValueError(f"Please specify the conversation mode for {model_path}.")
+                raise ValueError(f"Please specify the conversation mode for {output_dir}.")
             config["chat_template"] = conv_mode
             jinja_template = generate_jinja_template(conv_mode)
-            jinja_path = os.path.join(model_path, f"{conv_mode}.jinja")
+            jinja_path = os.path.join(output_dir, f"{conv_mode}.jinja")
             with open(jinja_path, "w") as f:
                 f.write(jinja_template)
         json.dump(config, open(cfg_path, "w"), indent=2)
-        self.copy_remote_py_files(model_path, copy=copy)
 
         ##########################################################################################
-        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-        tokenizer = load_tokenizer_then_handle_media_tokens_and_chat_template(model_path, config)
+        config = AutoConfig.from_pretrained(output_dir, trust_remote_code=True)
+        tokenizer = load_tokenizer_then_handle_media_tokens_and_chat_template(output_dir, config)
         tokenizer.save_pretrained(osp.join(output_dir, "llm"))
         ##########################################################################################
 
+    @classmethod
+    def copy_or_symlink_directory(cls, model_path, output_dir, copy=True):
+        # Create output directory if it doesn't exist
+        os.makedirs(output_dir, exist_ok=True)
+        # Create symlinks for all files in model_path to output_dir
+        for item in os.listdir(model_path):
+            src_path = os.path.join(model_path, item)
+            dst_path = os.path.join(output_dir, item)
+
+            # Remove existing file/directory at destination if it exists
+            if os.path.exists(dst_path):
+                if os.path.islink(dst_path):
+                    os.unlink(dst_path)
+                elif os.path.isdir(dst_path):
+                    shutil.rmtree(dst_path)
+                else:
+                    os.remove(dst_path)
+
+            # Create symlink
+            if copy:
+                if os.path.isdir(src_path):
+                    shutil.copytree(src_path, dst_path)
+                else:
+                    shutil.copy2(src_path, dst_path)
+                print(f"Copied {src_path} to {dst_path}")
+            else:
+                os.symlink(src_path, dst_path)
+                print(f"Created symlink from {src_path} to {dst_path}")
+
     @classmethod
     def copy_remote_py_files(cls, output_dir, copy=True):
         ## copy .py and REAMDE for next loading remote code
@@ -539,6 +588,15 @@ class VILAForCasualLM(VILAPretrainedModel):
             image_features = self.get_mm_projector()(image_features)
         return image_features
 
+    def train(self, mode: bool = True):
+        if mode:
+            print(f"Set padding side to right for training, {mode=}")
+            self.tokenizer.padding_side = "right"
+        else:
+            print(f"Set padding side to left for evaluation, {mode=}")
+            self.tokenizer.padding_side = "left"
+        super().train(mode)
+
     def _embed(
         self,
         input_ids: torch.Tensor,
@@ -547,18 +605,25 @@ class VILAForCasualLM(VILAPretrainedModel):
         labels: Optional[torch.Tensor],
         attention_mask: Optional[torch.Tensor],
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # NOTE(ligeng): deep copy to avoid modifying the original media and media_config
+        media = copy.deepcopy(media)
+        media_config = copy.deepcopy(media_config)
+
         labels = labels if labels is not None else torch.full_like(input_ids, IGNORE_INDEX)
         attention_mask = attention_mask if attention_mask is not None else torch.ones_like(input_ids, dtype=torch.bool)
 
-        # PROCESS_GROUP_MANAGER = get_pg_manager()
-        PROCESS_GROUP_MANAGER = None
+        PROCESS_GROUP_MANAGER = get_pg_manager()
         if PROCESS_GROUP_MANAGER is not None:
             for name in media:
                 self.encoders[name].end_tokens = None
 
         # Extract text and media embeddings
         text_embeds = self.llm.model.embed_tokens(input_ids)
-        media_embeds = self.__embed_media_tokens(media, media_config)
+        if media is not None:
+            media_embeds = self.__embed_media_tokens(media, media_config)
+        else:
+            # no media was provided, so we just return an empty dict
+            media_embeds = {}
 
         # This is a workaround to make sure the dummy embeddings are consumed
         while media_embeds.get("dummy"):
@@ -586,12 +651,20 @@ class VILAForCasualLM(VILAPretrainedModel):
                     name = media_tokens[input_ids[k][pos].item()]
                     input = media_embeds[name].popleft()
                     label = torch.full([input.shape[0]], IGNORE_INDEX, device=labels[k].device, dtype=labels[k].dtype)
+                    # print(f"{self.tokenizer.padding_side} [media] {k=} {pos=}, {self.tokenizer.batch_decode(input_ids[k][pos:pos+1])}"); python_input()
+                elif input_ids[k][pos].item() in (self.tokenizer.pad_token_id, self.tokenizer.eos_token_id):
+                    end = pos + 1
+                    pos = end
+                    # print(f"[skip PAD/EOS] {k=} {pos=}, {self.tokenizer.batch_decode(input_ids[k][pos:end])}"); python_input()
+                    continue
                 else:
                     end = pos
                     while end < len(labels[k]) and input_ids[k][end].item() not in media_tokens:
                         end += 1
                     input = text_embeds[k][pos:end]
                     label = labels[k][pos:end]
+                    # print(f"[text] {k=} {pos=}, {self.tokenizer.batch_decode(input_ids[k][pos:end])}"); python_input()
+
                 inputs_mk.append(input)
                 labels_mk.append(label)
                 pos = end
@@ -602,7 +675,7 @@ class VILAForCasualLM(VILAPretrainedModel):
         # Check if all media embeddings are consumed
         for name in media_embeds:
             if media_embeds[name]:
-                raise ValueError(f"Not all {name} embeddings are consumed!")
+                raise ValueError(f"Not all {name} embeddings are consumed! Still {len(media_embeds[name])} left.")
 
         # Truncate sequences to `model_max_length` as media embeddings are inserted
         inputs, labels = self.__truncate_sequence(inputs, labels)
@@ -620,7 +693,7 @@ class VILAForCasualLM(VILAPretrainedModel):
             if self.training:
                 # Gather metainfo of media objects from all ranks
                 info = [{"shape": tensor.shape, "dtype": tensor.dtype} for tensor in media.get(name, [])]
-                infos = list(chain(*distributed.all_gather(info)))
+                infos = list(chain(vila_all_gather(info)))
 
                 # The entire batch does not contain any media objects of this type.
                 if not infos:
@@ -1011,6 +1084,10 @@ class VILAForCasualLM(VILAPretrainedModel):
         attention_mask: Optional[torch.LongTensor] = None,
         **generation_kwargs,
     ):
+        if self.training:
+            warnings.warn(
+                "Model is in training mode, using default padding strategy to right. This is not recommended for generation."
+            )
         inputs_embeds, _, attention_mask = self._embed(input_ids, media, media_config, None, attention_mask)
         return self.llm.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, **generation_kwargs)
 

utils.py

@@ -41,7 +41,7 @@ def load_tokenizer_then_handle_media_tokens_and_chat_template(
         print(f"Using chat template: {config.chat_template}")
         fpath = os.path.join(os.path.dirname(__file__), "chat_templates", f"{config.chat_template}.jinja")
         if not os.path.exists(fpath):
-            fpath = os.path.join(os.path.dirname(model_name_or_path), f"{config.chat_template}.jinja")
+            fpath = os.path.join(model_name_or_path, f"{config.chat_template}.jinja")
         with open(fpath) as fd:
             chat_template = fd.read()
         tokenizer.chat_template = chat_template.replace("    ", "").replace("\n", "")

vision_tower/config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "runs/train/qwen25_2B_3x3-sft-20241118122815/model/vision_tower",
+  "_name_or_path": "NVILA-Lite-2B-hf-preview/vision_tower",
   "architectures": [
     "SiglipVisionModel"
   ],
@@ -17,7 +17,7 @@
   "patch_size": 14,
   "projection_dim": 2048,
   "projector_hidden_act": "gelu_fast",
-  "torch_dtype": "bfloat16",
+  "torch_dtype": "float16",
   "transformers_version": "4.46.0",
   "vision_use_head": false
 }

vision_tower/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:38c39a77befc808483de3c3a5d20f02af1e3111e15a3e9909f64b2ea7b553fd9
-size 826707904
+oid sha256:7ea53807bd8094c4acea02b1e4c9677207bcd635c44e9c6ec9a0862d558a26d0
+size 826707464