# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # 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 functools import json import os import time from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple import torch import torch.nn.functional as F from megatron.core import InferenceParams, ModelParallelConfig, parallel_state from safetensors.torch import load_file from torch.distributed.fsdp import FullStateDictConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import ShardingStrategy, StateDictType from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy from torch.nn.modules.module import _IncompatibleKeys from cosmos_predict1.autoregressive.configs.base.model import TrainingModelConfig as ModelConfig from cosmos_predict1.autoregressive.configs.base.tokenizer import TokenizerConfig from cosmos_predict1.autoregressive.modules.mm_projector import MultimodalProjector from cosmos_predict1.autoregressive.networks.vit import VisionTransformer, get_vit_config # from cosmos_predict1.autoregressive.training.networks.transformer_medusa import TransformerMedusa from cosmos_predict1.autoregressive.tokenizer.tokenizer import DiscreteMultimodalTokenizer from cosmos_predict1.autoregressive.training.networks.transformer import ( Transformer, TransformerBlock, TransformerBlockTE, ) from cosmos_predict1.autoregressive.utils.checkpoint import ( get_partial_state_dict, maybe_convert_checkpoint_to_backend, obtain_tensor_parallel_state_dict, process_state_dict, substrings_to_ignore, ) from cosmos_predict1.autoregressive.utils.misc import random_dropout from cosmos_predict1.autoregressive.utils.parallel import broadcast_data_batch_in_tp_cp_group, get_batch_on_this_cp_rank from cosmos_predict1.autoregressive.utils.sampling import ( decode_n_tokens, decode_one_token, prefill, sample_top_k, sample_top_p, ) from cosmos_predict1.diffusion.training.utils.fsdp_helper import apply_fsdp_checkpointing, hsdp_device_mesh from cosmos_predict1.utils import distributed, log, misc from cosmos_predict1.utils.lazy_config import LazyDict from cosmos_predict1.utils.misc import download_from_s3_with_cache, sync_s3_dir_to_local from cosmos_predict1.utils.model import Model class AutoRegressiveTrainingModel(Model): """ A class to build and use a Llama model for text generation. Methods: build: Build a Llama instance by initializing and loading a model checkpoint. generate: Generate text sequences based on provided prompts using the language generation model. """ def __init__( self, model: Transformer, tokenizer: DiscreteMultimodalTokenizer, config: ModelConfig, model_parallel: ModelParallelConfig = None, vision_encoder: VisionTransformer = None, mm_projector: MultimodalProjector = None, ): """ Initialize the Llama instance with a model and tokenizer. Args: model (Transformer): The Transformer model for text generation. tokenizer (Tokenizer): The tokenizer for encoding and decoding text. config (Config): The configuration for the Llama model. """ super().__init__() self.model = model self.tokenizer = tokenizer self.config = config self.precision = self.model.precision self.vision_encoder = vision_encoder self.mm_projector = mm_projector assert (self.vision_encoder is None) == (self.mm_projector is None), ( "vision_encoder and mm_projector should be " "both None or not None simultaneously" ) self.model_parallel = model_parallel self.monitor_output_logits = False self.inference_params = None # self.insert_medusa_head = self.config.insert_medusa_head if self.config.freeze_vision_encoder and vision_encoder is not None: for param in self.vision_encoder.parameters(): param.requires_grad = False log.critical("Vision encoder parameters are frozen.") num_params = self.get_num_params() log.info(f"Number of model parameters: {round(num_params / 1e9, 3)}B") def get_num_params( self, ) -> int: """ Return the number of parameters in the model. """ n_params = sum(p.numel() for p in self.parameters()) return n_params def training_step( self, data_batch: dict[str, torch.Tensor], iteration: int ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: broadcast_data_batch_in_tp_cp_group(data_batch) # get the context embedding and mask context = data_batch.get("context", None) context_mask = data_batch.get("context_mask", None) if context is not None: if self.config.embedding_dropout > 0: context = random_dropout( context, self.config.embedding_dropout, ) context = misc.to(context, device="cuda") if context_mask is not None: context_mask = misc.to(context_mask, device="cuda") action = data_batch.get("action", None) if action is not None: action = misc.to(action, device="cuda") # Input tokens tokens, token_boundaries = self.tokenizer.tokenize(data_batch) tokens = misc.to(tokens, device="cuda") # Tokens to predict labels = data_batch.get("labels", None) # Token Mask (Note: this is not attention mask) masks = data_batch.get("token_mask", None) apply_token_mask = masks is not None if masks is None: masks = torch.ones_like(tokens, dtype=torch.bool) masks = misc.to(masks, device="cuda") assert ( data_batch.get("labels", None) is None or apply_token_mask ), "The code is not tested for the case when both labels and token_mask are provided." if self.config.ignore_first_num_tokens > 0: assert self.config.ignore_first_num_tokens < masks.shape[1] masks[:, : self.config.ignore_first_num_tokens] = False seq_len = tokens.shape[1] # Boradcast inputs to TP and CP ranks, alternatively we can use the `_broadcast` function from cosmos/diffusion/v1 # Currently we only handled video tokens (with label and mask) and text tokens (with mask), action and other inputs might also need to be handled if parallel_state.get_context_parallel_world_size() > 1: # Turn on CP cp_group = parallel_state.get_context_parallel_group() self.model.enable_context_parallel(cp_group) tokens = get_batch_on_this_cp_rank(tokens) masks = get_batch_on_this_cp_rank(masks) if labels is not None: labels = get_batch_on_this_cp_rank(labels) if self.vision_encoder is None: logits = self.model.forward( tokens=tokens, input_pos=None, context=context, context_mask=context_mask, action=action, total_seq_len=seq_len, ) else: assert "images" in data_batch images = data_batch["images"] if images.ndim == 5: # The shape is (batch_size, n_images_per_sample, C, H, W). Flatten the first two dimensions. images = images.view(-1, *images.shape[2:]) assert images.ndim == 4, f"Invalid shape: {images.shape}" token_embeddings = self.embed_vision_language_features(tokens, images) logits = self.model.forward( token_embeddings=token_embeddings, input_pos=None, context=context, context_mask=context_mask, action=action, total_seq_len=seq_len, ) if labels is None: # For auto-regressive models, the labels are the same as the # input tokens shifted by one position logits = logits[:, :-1] masks = masks[:, :-1] labels = tokens[:, 1:].clone() batch_size = tokens.shape[0] # Apply ignore_index for sample_num in range(batch_size): if self.tokenizer.training_type == "text_to_video": # For text-to-video training, we do not compute the loss of text part # Hence, we set the labels of text tokens to that of ignore_index if len(token_boundaries["text"]) > 0: labels[sample_num][0 : token_boundaries["text"][sample_num][1] - 1] = self.tokenizer.ignore_index elif self.tokenizer.training_type == "class_to_image": # For class-to-image training, we do not compute the loss of class part # Hence, we set the labels of class tokens to that of ignore_index labels[sample_num][0 : token_boundaries["class"][sample_num][1] - 1] = self.tokenizer.ignore_index ignore_index = self.tokenizer.ignore_index if self.config.ignore_first_num_tokens > 0 or apply_token_mask: labels[~masks] = ignore_index output_batch = { "encode_tokens": tokens, "logits": logits.detach(), "labels": labels.detach(), "ignore_index": ignore_index, } if self.monitor_output_logits: self.gather_output_logits_stats(logits, labels, output_batch, ignore_index) logits = logits.flatten(0, 1) labels = labels.flatten(0, 1) # Main cross entropy loss ce_loss = F.cross_entropy( input=logits, target=labels, ignore_index=ignore_index, # ignore prompt (turn prompt tokens into pad_id here) ) # Z-loss log_z = torch.logsumexp(logits, dim=-1) # shape: [B, seq_len] z_loss = self.config.z_loss_coeff * (log_z**2).mean() # Combined loss total_loss = ce_loss + z_loss return output_batch, total_loss # skip returning output logits @torch.no_grad() def validation_step( self, data_batch: dict[str, torch.Tensor], iteration: int ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: """ Perform a validation step for the model, which is the same as the training step (but without backpropagation). """ return self.training_step(data_batch, iteration) @torch.no_grad() def gather_output_logits_stats( self, logits: torch.Tensor, labels: torch.Tensor, output_batch: Dict, ignore_index: int = None ): """ Gather statistics of the output logits, including mean, norm, and max values. """ bs, seq_len, dim = logits.shape logits = logits.reshape(-1, dim) if ignore_index is not None: select_index = labels.view(-1) != ignore_index acc = labels.view(-1)[select_index] == logits.argmax(dim=1)[select_index] acc = acc.float().mean().view(-1, 1) logits = logits[select_index] output_batch.update( { "logits_mean": logits.mean(dim=1).detach(), "logits_norm": torch.linalg.vector_norm(logits, dim=1).detach(), "logits_max": logits.max(dim=1).values.detach(), "acc": acc.detach() * 100, } ) @torch.no_grad() def image_encode(self, state: torch.Tensor) -> torch.Tensor: """ Encode the image input state to continuous latent and discrete indices. """ latent, indices = self.tokenizer.image_tokenizer.encode(state) return latent, indices @torch.no_grad() def image_decode(self, indices: torch.Tensor) -> torch.Tensor: """ Decode the discrete indices to RGB images. """ return self.tokenizer.image_tokenizer.decode(indices) @torch.no_grad() def video_encode(self, state: torch.Tensor) -> torch.Tensor: """ Encode the video input state to continuous latent and discrete indices. """ latent, indices = self.tokenizer.video_tokenizer.encode(state) return latent, indices @torch.no_grad() def video_decode(self, indices: torch.Tensor) -> torch.Tensor: """ Decode the discrete indices to RGB videos. """ if self.tokenizer.tokenizer_config.video_tokenizer.temporal_overlap > 0: return self.tokenizer.video_tokenizer.decode_with_overlap( indices, temporal_overlap=self.tokenizer.tokenizer_config.video_tokenizer.temporal_overlap ) else: return self.tokenizer.video_tokenizer.decode(indices) @staticmethod def load_llm_checkpoint( ckpt_path: str = "", model: Transformer = None, **kwargs, ) -> None: """ Load a LLM checkpoint from the specified path. """ with misc.timer(f"loading checkpoint from {ckpt_path}"): checkpoint = torch.load( ckpt_path, map_location="cpu", mmap=True, # load the checkpoint in memory-mapped mode ) llm_checkpoint = checkpoint["model"] if "model" in checkpoint else checkpoint llm_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="model.") llm_checkpoint = process_state_dict(llm_checkpoint, prefix_to_remove="model.") with misc.timer("loading state_dict into model"): missing_keys, unexpected_keys = model.load_state_dict(llm_checkpoint, strict=True) @staticmethod def build( seed: int = 1, train_from_scratch: bool = False, model_config: ModelConfig = ModelConfig(), fsdp_checkpointer: Any = None, tokenizer_config: TokenizerConfig = None, model_parallel: ModelParallelConfig = None, shard_checkpoint: bool = True, download_rank_sync: bool = True, **kwargs, ) -> "AutoRegressiveTrainingModel": """ Build a Llama instance by initializing and loading a model checkpoint. Args: seed (int, optional): Random seed for reproducibility. Defaults to 1. train_from_scratch (bool, optional): Flag indicating whether to train the model from scratch. Defaults to False. model_config (ModelConfig, optional): The model configuration for the Llama instance. Defaults to ModelConfig(). fsdp_checkpointer (Any, optional): The FSDP checkpointer for the Llama instance. Defaults to None. tokenizer_config (TokenizerConfig, optional): The tokenizer configuration for the Llama instance. Defaults to None. shard_checkpoint (bool, optional): Whether to split the checkpoint by Tensor Parallelism before loading. Defaults to False. download_rank_sync (bool, optional): Whether to download the checkpoint in a rank-synchronized manner. Defaults to True. Returns: Llama: An instance of the Llama class with the loaded model and tokenizer. Raises: AssertionError: If there are no checkpoint files in the specified directory. Note: This method sets the device to CUDA and loads the pre-trained model and tokenizer. """ tensor_parallel_size = 1 if model_parallel is None else model_parallel.tensor_model_parallel_size # seed must be the same in all processes torch.manual_seed(seed) # Initialize model configuration parameters llama_params = {} # Load checkpoint and model parameters if not train_from_scratch: if model_config.ckpt_path is None: # If ckpt_path is not provided, we assume the model checkpoint is saved in the ckpt_dir ckpt_dir = sync_s3_dir_to_local( s3_dir=model_config.ckpt_dir, s3_credential_path=model_config.s3_credential_path, cache_dir=model_config.cache_dir, ) # We prioritize safetensors version over the pytorch version, since the former is # much faster for checkpoint loading. checkpoints = sorted(Path(ckpt_dir).glob("*.safetensors")) if len(checkpoints) == 0: checkpoints = sorted(Path(ckpt_dir).glob("*.pth")) assert len(checkpoints) > 0, f"no checkpoint files found in {ckpt_dir}" assert ( len(checkpoints) == 1 ), f"multiple checkpoint files found in {ckpt_dir} (currently only one is supported)" ckpt_path = str(checkpoints[0]) # Assuming single checkpoint for non-parallel case if os.path.exists(Path(ckpt_dir) / "params.json"): with open(Path(ckpt_dir) / "params.json", "r") as f: llama_params = json.loads(f.read()) else: log.info( f"No params.json found in the checkpoint directory ({ckpt_dir}). " f"Using default model config." ) else: # If ckpt_path is provided, we load the model from the specified path, # and use the default model configuration ckpt_path = download_from_s3_with_cache( s3_path=model_config.ckpt_path, s3_credential_path=model_config.s3_credential_path, cache_dir=model_config.cache_dir, rank_sync=download_rank_sync, ) for key, value in llama_params.items(): # Override the default model configuration with the parameters from the checkpoint setattr(model_config, key, value) with misc.timer(f"loading checkpoint from {ckpt_path}"): if ckpt_path.endswith("safetensors"): # Load with safetensors API checkpoint = load_file(ckpt_path, device="cpu") else: # The pytorch version checkpoint = torch.load( ckpt_path, map_location="cpu", mmap=True, # load the checkpoint in memory-mapped mode ) llm_checkpoint = checkpoint["model"] if "model" in checkpoint else checkpoint # If the checkpoint backend is different from the model backend, convert the checkpoint # to be compatible with the model backend # If shard_checkpoint is True, the loaded checkpoint is the whole model checkpoint (will be sharded later) # instead of a tensor-parallel sharded checkpoint llm_checkpoint = maybe_convert_checkpoint_to_backend( llm_checkpoint, target_backend=model_config.backend, model_config=model_config, tensor_parallel_size=tensor_parallel_size if not shard_checkpoint else 1, is_tensor_parallel_shard=tensor_parallel_size > 1 and not shard_checkpoint, ) if model_config.vision_encoder is not None: # For vanilla VLM ckpt before fine-tuning, `checkpoint['model']` only contains LLM weights, and `checkpoint['vision_encoder']` # and `checkpoint['mm_projector']` are both for those weights # For fine-tuned VLM ckpt, `checkpoint['model']` contains all LLM, mm_projector and vision_encoder weights if "vision_encoder" in checkpoint: log.info("Using pretrained vision_encoder") vit_checkpoint = checkpoint["vision_encoder"] else: log.info("Using fine-tuned vision_encoder") vit_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="vision_encoder.") vit_checkpoint = process_state_dict(vit_checkpoint, prefix_to_remove="vision_encoder.") if "mm_projector" in checkpoint: log.info("Using pretrained mm_projector") projector_checkpoint = checkpoint["mm_projector"] else: log.info("Using fine-tuned mm_projector") projector_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="mm_projector.") projector_checkpoint = process_state_dict(projector_checkpoint, prefix_to_remove="mm_projector.") assert ( len(vit_checkpoint) > 0 and len(projector_checkpoint) > 0 ), "vit_checkpoint and projector_checkpoint cannot be empty. We do not support random initialization for vision_encoder and mm_projector." tokenizer = DiscreteMultimodalTokenizer(tokenizer_config) precision = getattr(torch, model_config.precision) torch.set_default_dtype(precision) log.info(f"Setting torch default dtype to {precision}") # if model_config.insert_medusa_head: # model = TransformerMedusa( # params=model_config, # model_parallel=model_parallel, # tokenizer_config=tokenizer_config, # init_weights=train_from_scratch, # ) # else: model = Transformer( params=model_config, model_parallel=model_parallel, tokenizer_config=tokenizer_config, init_weights=train_from_scratch, ) model_kwargs = {} # [Optional] Initialize vision encoder and multimodal projector (for vision-language tasks) if model_config.vision_encoder is not None: assert model_config.mm_projector is not None, "mm_projector must be provided if vision_encoder is provided." vit_config = get_vit_config(model_config.vision_encoder) vision_encoder = VisionTransformer.build( vit_config, hidden_dropout=model_config["hidden_dropout"], attention_dropout=model_config["attention_dropout"], set_parallel_mode=model_config["set_parallel_mode"], model_parallel=model_parallel, attention_tp=tensor_parallel_size > 1, ) mm_projector = MultimodalProjector( mm_projector_type=model_config.mm_projector, in_dim=vit_config["dim"], out_dim=model_config["dim"] ) model_kwargs.update({"vision_encoder": vision_encoder, "mm_projector": mm_projector}) # Perform vocab expansion if tokenizer.vocab_size > model.vocab_size: log.info(f"Expanding vocab size to {tokenizer.vocab_size}") # For text-to-video training, we only expand the embedding layer but not the output (unembedding) layer, expand_output_layer = not (tokenizer.training_type == "text_to_video") model.expand_vocab(tokenizer.vocab_size, init_method="gaussian", expand_output_layer=expand_output_layer) if not train_from_scratch: if shard_checkpoint: # Shard the checkpoint according to tensor parallelism. with misc.timer("sharding checkpoint according to tensor parallelism"): if model_parallel is not None: assert model_parallel.tensor_model_parallel_size == model_config["tensor_model_parallel_size"] llm_checkpoint = obtain_tensor_parallel_state_dict( llm_checkpoint, tensor_parallel_size=tensor_parallel_size, tensor_parallel_rank=parallel_state.get_tensor_model_parallel_rank(), model_config=model_config, ) if model_config.vision_encoder is not None: # Shard vision encoder and multimodal projector weights vit_checkpoint = obtain_tensor_parallel_state_dict( vit_checkpoint, tensor_parallel_size=tensor_parallel_size, tensor_parallel_rank=parallel_state.get_tensor_model_parallel_rank(), model_config=vit_config, ) if model_config.vision_encoder is not None: # Take the LLM weights (starting with "model.") from the VLM checkpoint llm_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="model.") # Remove the "model." prefix in the state_dict llm_checkpoint = process_state_dict(llm_checkpoint, prefix_to_remove="model.") with misc.timer("loading state_dict into model"): missing_keys, unexpected_keys = model.load_state_dict(llm_checkpoint, strict=True) # Remove keys with "_extra_state" suffix in missing_keys (defined by TransformerEngine for FP8 usage) missing_keys = [k for k in missing_keys if not k.endswith("_extra_state")] assert len(missing_keys) == 0, f"Missing keys: {missing_keys}" if model_config.vision_encoder is not None: # Load vision encoder and multimodal projector weights vision_encoder.load_state_dict(vit_checkpoint) mm_projector.load_state_dict(projector_checkpoint) if model_config.vision_encoder_in_channels != 3: vision_encoder.expand_in_channels(model_config.vision_encoder_in_channels) model = model.to(precision) # ensure model parameters are in the correct precision log.info(f"Model config: {model_config}") # if model_config.insert_medusa_head: # from projects.cosmos.ar.v1.model_medusa import LlamaMedusa # model_class = LlamaMedusa # else: model_class = AutoRegressiveTrainingModel if model_config.fsdp_enabled: raise NotImplementedError("FSDP is not implemented for AutoRegressiveTrainingModel") # model_kwargs["fsdp_checkpointer"] = fsdp_checkpointer # model_class = FSDPLlama return model_class(model, tokenizer, model_config, **model_kwargs) @torch.inference_mode() def generate( self, prompt_tokens: List[List[int]], max_gen_len: int, temperature: float = 0.6, top_p: float = 0.9, top_k: Optional[int] = None, logprobs: bool = False, echo: bool = False, logit_clipping_range: list = [], seed: int = 0, images: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, action: Optional[torch.Tensor] = None, ) -> Tuple[List[List[int]], Optional[List[List[float]]]]: """ Generate text sequences based on provided prompts using the language generation model. Args: prompt_tokens (List[List[int]]): List of tokenized prompts, where each prompt is represented as a list of integers. max_gen_len (int): Maximum length of the generated text sequence. temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6. top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to 0.9. top_k (int, optional): Top-k value for top-k sampling. Defaults to None. If not None, top-k sampling will be used instead of top-p sampling. logprobs (bool, optional): Flag indicating whether to compute token log probabilities. Defaults to False. echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False. Returns: Tuple[List[List[int]], Optional[List[List[float]]]]: A tuple containing generated token sequences and, if logprobs is True, corresponding token log probabilities. Note: This method uses the provided prompts as a basis for generating text. It employs nucleus sampling to produce text with controlled randomness. If logprobs is True, token log probabilities are computed for each generated token. """ assert top_k is None or top_p is None, f"Only one of top_k ({top_k} or top_p ({top_p} should be specified." if top_p is not None: log.info(f"Using top-p sampling with p={top_p} and temperature={temperature}") elif top_k is not None: log.info(f"Using top-k sampling with k={top_k} and temperature={temperature}") else: log.info("Not applying top-k or top-p sampling. Will use top-k sampling with k=None") self.model.set_inference_flag(True) misc.set_random_seed(seed) # Initialization and Assertions if isinstance(self.model.params, list): # During training, model.params is a list log.info( f"Find self.model.params is a list, use self.config instead. Get max_batch_size={self.config.max_batch_size}, max_seq_len={self.config.max_seq_len}" ) params = self.config else: params = self.model.params bsz = len(prompt_tokens) assert bsz <= params.max_batch_size, (bsz, params.max_batch_size) if self.config.backend == "transformer_engine": self.inference_params = InferenceParams( max_batch_size=params.max_batch_size, max_sequence_length=params.max_seq_len ) # Calculate Prompt Lengths min_prompt_len = min(len(t) for t in prompt_tokens) max_prompt_len = max(len(t) for t in prompt_tokens) assert max_prompt_len <= params.max_seq_len total_len = params.max_seq_len assert ( max_gen_len + max_prompt_len <= total_len ), f"max_gen_len + max_prompt_len={max_gen_len + max_prompt_len} exceeds max_seq_len={total_len}" pad_id = self.tokenizer.pad_id tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device="cuda") # Fill tokens tensor with prompt tokens for k, t in enumerate(prompt_tokens): tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device="cuda") if logprobs: token_logprobs = torch.zeros_like(tokens, dtype=torch.float) prev_pos = 0 eos_reached = torch.tensor([False] * bsz, device="cuda") input_text_mask = tokens != pad_id # Flag to check if image embeddings have been passed to the model - we only need to pass them once # since we have KV cache. passed_image_embeddings = False # If all prompts are of max length, compute initial logits and logprobs if min_prompt_len == total_len: input_pos = torch.arange(tokens.shape[1], dtype=torch.long, device="cuda") if images is None: logits = self.model.forward( tokens=tokens, input_pos=input_pos, inference_params=self.inference_params, context=context, context_mask=context_mask, action=action, ) else: token_embeddings = self.embed_vision_language_features(tokens, images) logits = self.model.forward( token_embeddings=token_embeddings, input_pos=input_pos, inference_params=self.inference_params, context=context, context_mask=context_mask, action=action, ) passed_image_embeddings = True token_logprobs = -F.cross_entropy( input=logits.transpose(1, 2), target=tokens, reduction="none", ignore_index=pad_id, ) stop_tokens = torch.tensor(list(self.tokenizer.stop_tokens), dtype=torch.long, device="cuda") # Main generation loop log.info(f"Start generating the next {total_len - min_prompt_len} tokens. This will take a while..") for cur_pos in range(min_prompt_len, total_len): input_pos = torch.arange(prev_pos, cur_pos, dtype=torch.long, device="cuda") if images is not None and not passed_image_embeddings: token_embeddings = self.embed_vision_language_features(tokens[:, prev_pos:cur_pos], images) logits = self.model.forward( token_embeddings=token_embeddings, input_pos=input_pos, inference_params=self.inference_params, context=context, context_mask=context_mask, action=action, ) passed_image_embeddings = True else: logits = self.model.forward( tokens=tokens[:, prev_pos:cur_pos], input_pos=input_pos, inference_params=self.inference_params, context=context, context_mask=context_mask, action=action, ) if self.config.backend == "transformer_engine": self.inference_params.sequence_len_offset += logits.shape[1] # Apply temperature scaling and nucleus sampling if len(logit_clipping_range) > 0: min_clip_index = logit_clipping_range[0] max_clip_index = logit_clipping_range[1] logits_clipped = logits[:, :, min_clip_index:max_clip_index] else: logits_clipped = logits min_clip_index = 0 if temperature > 0: if top_p is not None: next_token = sample_top_p(logits_clipped, temperature=temperature, top_p=top_p)[0] else: next_token = sample_top_k(logits_clipped, temperature=temperature, top_k=top_k)[0] else: next_token = torch.argmax(logits_clipped[:, -1, :], dim=-1) next_token += min_clip_index next_token = next_token.reshape(-1) # only replace token if prompt has already been generated next_token = torch.where(input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token) tokens[:, cur_pos] = next_token # Calculate log probabilities if requested if logprobs: token_logprobs[:, prev_pos + 1 : cur_pos + 1] = -F.cross_entropy( input=logits.transpose(1, 2), target=tokens[:, prev_pos + 1 : cur_pos + 1], reduction="none", ignore_index=pad_id, ) # Check if end-of-sequence token is reached eos_reached |= (~input_text_mask[:, cur_pos]) & (torch.isin(next_token, stop_tokens)) prev_pos = cur_pos # Break the loop if all sequences have reached an end-of-sequence token if all(eos_reached): log.info(f"Reach end of sequence, current pos: {cur_pos}; maximum pos: {total_len}") break # Convert log probabilities to list if required if logprobs: token_logprobs = token_logprobs.tolist() out_tokens, out_logprobs = [], [] # Process and collect the output tokens and log probabilities for i, toks in enumerate(tokens.tolist()): # cut to max gen len start = 0 if echo else len(prompt_tokens[i]) toks = toks[start : len(prompt_tokens[i]) + max_gen_len] probs = None if logprobs: probs = token_logprobs[i][start : len(prompt_tokens[i]) + max_gen_len] # cut to after eos tok if any for stop_token in self.tokenizer.stop_tokens: try: eos_idx = toks.index(stop_token) toks = toks[:eos_idx] probs = probs[:eos_idx] if logprobs else None except ValueError: pass out_tokens.append(toks) out_logprobs.append(probs) self.model.set_inference_flag(False) return (out_tokens, out_logprobs if logprobs else None) @torch.no_grad() def fast_generate( self, prompt_tokens: List[List[int]] | torch.Tensor, max_gen_len: int, temperature: float = 1.0, top_k: Optional[int] = None, top_p: Optional[float] = None, num_gen_seq: int = 1, logprobs: bool = False, echo: bool = False, seed: int = 0, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, action: Optional[torch.Tensor] = None, compile_decode: bool = True, compile_prefill: bool = False, verbose: bool = True, stop_tokens: Optional[Set[int]] = None, ): """ Fast auto-regressive generation. Currently only supports input batch size = 1. Args: prompt_tokens (List[List[int]] | torch.Tensor): A single prompt of shape (1, seq_len). max_gen_len (int): Maximum length of the generated text sequence. temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6. top_k (int, optional): Top-k value for top-k sampling. Defaults to None. top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to None. num_gen_seq (int, optional): Number of outputs to generate given the same prompt. Defaults to 1. When temperature == 0, num_gen_seq must be 1 because the generation is deterministic. echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False. logit_clipping_range (list, optional): Range of logits to clip. Defaults to []. seed (int, optional): Random seed for reproducibility. Defaults to 0. compile_decode (bool, optional): Flag indicating whether to compile the decoding function. Defaults to True. compile_prefill (bool, optional): Flag indicating whether to compile the prefill function. Defaults to False. verbose (bool, optional): Flag indicating whether to print the the time. Defaults to False. """ assert ( top_p is None or top_k is None ), f"Only one of top-p or top-k can be provided, got top-p={top_p} and top-k={top_k}" if top_p is not None: log.info(f"Using top-p sampling with p={top_p} and temperature={temperature}") elif top_k is not None: log.info(f"Using top-k sampling with k={top_k} and temperature={temperature}") else: log.info("Not applying top-k or top-p sampling. Will use top-k sampling with k=None") torch._inductor.config.coordinate_descent_tuning = True torch._inductor.config.triton.unique_kernel_names = True # Experimental features to reduce compilation times, will be on by default in future torch._inductor.config.fx_graph_cache = True # torch._functorch.config.enable_autograd_cache = True self.model.set_inference_flag(True) misc.set_random_seed(seed) assert not logprobs, "logprobs are not supported for fast_generate yet" # Examine if the function prefil and decode_one_token functions are compiled yet. If not, compile them based on the flags if compile_decode and not getattr(self, "inference_decode_compiled", False): self.decode_one_token = torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True) self.inference_decode_compiled = True log.critical("Compiled decode_one_token function. Note: the first run will be slower due to compilation") if compile_prefill and not getattr(self, "inference_prefill_compiled", False): self.prefill = torch.compile(prefill, fullgraph=True, dynamic=True) self.inference_prefill_compiled = True log.critical("Compiled prefill function. Note: the first run will be slower due to compilation") if not hasattr(self, "decode_one_token"): self.decode_one_token = decode_one_token if not hasattr(self, "prefill"): self.prefill = prefill # Initialization and Assertions if isinstance(self.model.params, list): # During training, model.params is a list log.info( f"Find self.model.params is a list, use self.config instead. Get max_batch_size={self.config.max_batch_size}, max_seq_len={self.config.max_seq_len}" ) params = self.config else: params = self.model.params if isinstance(prompt_tokens, list): prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device="cuda") if prompt_tokens.ndim == 1: prompt_tokens = prompt_tokens.view(1, -1) else: assert prompt_tokens.ndim == 2, f"prompt_tokens has shape {prompt_tokens.shape}" batch_size, prompt_len = prompt_tokens.shape total_len = min(params.max_seq_len, max_gen_len + prompt_len) if max_gen_len + prompt_len > params.max_seq_len: log.warning( f"max_gen_len + prompt_len={max_gen_len + prompt_len} exceeds max_seq_len={params.max_seq_len}, truncate max_gen_len to {params.max_seq_len - prompt_len}" ) max_gen_len = params.max_seq_len - prompt_len if context_mask is not None: context_mask = context_mask.to(dtype=torch.bool) if context_mask.ndim == 2: assert ( context_mask.shape[0] == batch_size ), f"batch_size mismatch: {context_mask.shape[0]} != {batch_size}" # Unsqueeze it to make it of shape [batch_size, 1, 1, context_seq_len] context_mask = context_mask.view(batch_size, 1, 1, -1) if num_gen_seq > 1: assert ( batch_size == 1 ), f"num_gen_seq > 1 is only supported for a single prompt, got {len(prompt_tokens)} prompts" log.critical(f"Generating {num_gen_seq} sequences with the same prompt") assert ( num_gen_seq <= params.max_batch_size ), f"num_gen_seq={num_gen_seq} exceeds max_batch_size={params.max_batch_size}" # repeat the prompt tokens for num_gen_seq times prompt_tokens = prompt_tokens.repeat(num_gen_seq, 1) assert prompt_tokens.shape == ( num_gen_seq, prompt_len, ), f"prompt_tokens must be of shape (num_gen_seq, seq_len), got {prompt_tokens.shape}" batch_size = len(prompt_tokens) # create an empty tensor of the expected final shape and fill in the current tokens empty = torch.empty(batch_size, total_len, dtype=prompt_tokens.dtype, device=prompt_tokens.device) empty[:, :prompt_len] = prompt_tokens seq = empty input_pos = torch.arange(0, prompt_len, device="cuda") if verbose: prefill_start = time.time() # Prefill stage next_token = self.prefill( self.model, prompt_tokens, input_pos=input_pos, temperature=temperature, top_k=top_k, top_p=top_p, context=context, context_mask=context_mask, action=action, ) if verbose: prefill_time = time.time() - prefill_start seq[:, [prompt_len]] = next_token.to(dtype=seq.dtype) input_pos = torch.tensor([prompt_len], dtype=torch.long, device="cuda") stop_tokens = self.tokenizer.stop_tokens if stop_tokens is None else stop_tokens stop_tokens = torch.tensor(list(stop_tokens), dtype=torch.long, device="cuda") if verbose: decode_start = time.time() # Decode stage generated_tokens = decode_n_tokens( self.model, next_token.view(batch_size, -1), input_pos, max_gen_len - 1, temperature=temperature, top_k=top_k, top_p=top_p, stop_tokens=stop_tokens, decode_one_token_function=self.decode_one_token, context=context, context_mask=context_mask, action=action, ) gen_len = len(generated_tokens) if verbose: decode_time = time.time() - decode_start prefill_throughput = prompt_len / prefill_time decode_throughput = gen_len / decode_time log.info(f"[Prefill] Time: {prefill_time:.2f}s; Throughput: {prefill_throughput:.2f} tokens/s") log.info(f"[Decode] Time: {decode_time:.2f}s; Throughput: {decode_throughput:.2f} tokens/s") generated_tokens = torch.cat(generated_tokens, dim=1) log.critical(f"generated_tokens: {generated_tokens.shape}") seq = seq[:, : prompt_len + 1 + gen_len] seq[:, prompt_len + 1 :] = generated_tokens if not echo: seq = seq[:, prompt_len:] return seq, None def embed_vision_language_features(self, input_ids: torch.Tensor, images: torch.tensor) -> torch.Tensor: """ Embed vision and language features into a combined representation. Args: input_ids (torch.Tensor): Input token IDs. images (torch.tensor): Input images. Returns: torch.Tensor: Combined vision-language features. Raises: AssertionError: If vision encoder or mm projector is not initialized, or if dimensions mismatch. """ # Ensure vision encoder and mm projector are initialized assert self.vision_encoder is not None assert self.mm_projector is not None # Get image token ID and validate it image_token_id = self.vision_encoder.image_token_id assert isinstance(image_token_id, int) and image_token_id >= 0, f"Invalid image_token_id: {image_token_id}" # Identify text and image locations in the input text_locations = input_ids != image_token_id image_locations = input_ids == image_token_id # Process text features text_features = self.model.tok_embeddings(input_ids[text_locations]) # Process image features images = images.to(device=text_features.device, dtype=text_features.dtype) vit_outputs = self.vision_encoder(images) image_features = self.mm_projector(vit_outputs) # Get dimensions B, seq_len = input_ids.shape N_total = B * seq_len N_txt, D_txt = text_features.shape N_img, N_patch, D_img = image_features.shape # Reshape image features image_features = image_features.reshape(N_img * N_patch, D_img) # Validate dimensions assert D_txt == D_img, f"Text features dim {D_txt} should be equal to image features dim {D_img}" assert ( N_total == N_txt + N_img * N_patch ), f"seq_len {seq_len} should be equal to N_txt + N_img*N_Patch {(N_txt, N_img * N_patch, image_locations.sum().item())}" # Combine text and image features combined_features = torch.empty( (B, seq_len, D_txt), dtype=text_features.dtype, device=text_features.device, ) combined_features[text_locations, :] = text_features combined_features[image_locations, :] = image_features return combined_features def on_after_backward(self, iteration: int = 0): """ Hook after loss.backward() is called. This method is called immediately after the backward pass, allowing for custom operations or modifications to be performed on the gradients before the optimizer step. So far, this method is used to all-reduce layernorm grads for tensor/sequence parallelism. Args: iteration (int): Current iteration number. """ for module in self.children(): if hasattr(module, "on_after_backward"): module.on_after_backward(iteration) def on_before_zero_grad( self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int ) -> None: """Hook before zero_grad() is called. Args: optimizer (torch.optim.Optimizer): The model optimizer. scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler. iteration (int): Current iteration number. """ for module in self.children(): if hasattr(module, "on_before_zero_grad"): module.on_before_zero_grad(optimizer, scheduler, iteration) @property def fsdp_wrap_block_cls(self): """ Return the transformer block class to wrap with FSDP. """ if self.config.backend == "pytorch": return TransformerBlock elif self.config.backend == "transformer_engine": return TransformerBlockTE else: raise ValueError(f"Unknown backend: {self.config.backend}") def state_dict(self, *args, **kwargs): """ Process the state dict (e.g., remove "_extra_state" keys imposed by TransformerEngine for FP8). """ state_dict = super().state_dict(*args, **kwargs) return process_state_dict(state_dict) def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True, assign: bool = False): """ Ignore the missing keys with substrings matching `substring_to_ignore` (e.g., "_extra_state" keys imposed by TransformerEngine for FP8). """ state_dict = process_state_dict(state_dict) missing_keys, unexpected_keys = super().load_state_dict(state_dict, strict=False, assign=assign) actual_missing_keys = [] for key in missing_keys: if not any(substring in key for substring in substrings_to_ignore): actual_missing_keys.append(key) if strict: if len(actual_missing_keys) > 0 or len(unexpected_keys) > 0: raise ValueError(f"Missing keys: {actual_missing_keys}\n\nUnexpected keys: {unexpected_keys}") return _IncompatibleKeys(actual_missing_keys, unexpected_keys) # class FSDPLlama(Llama): # def __init__( # self, model: Transformer, tokenizer: DiscreteMultimodalTokenizer, config: ModelConfig, fsdp_checkpointer: Any # ): # self.fsdp_checkpointer = fsdp_checkpointer # super().__init__(model, tokenizer, config) # self.set_up_fsdp() # def set_up_fsdp(self): # """ # Set up FSDP for the model. # """ # model = self.model # # detach the model from the parent class # self.model = None # del self.model # # build FSDP sharding strategy and device_mesh # strategy = { # "full": ShardingStrategy.FULL_SHARD, # "hybrid": ShardingStrategy.HYBRID_SHARD, # "none": ShardingStrategy.NO_SHARD, # }[self.config.fsdp["sharding_strategy"]] # log.critical(f"Using {strategy} sharding strategy for FSDP") # if self.config.fsdp["sharding_strategy"] == "hybrid": # sharding_group_size = self.config.fsdp["sharding_group_size"] # device_mesh = hsdp_device_mesh( # sharding_group_size=sharding_group_size, # ) # else: # device_mesh = hsdp_device_mesh( # sharding_group_size=distributed.get_world_size(), # ) # parallel_state.fsdp_device_mesh = device_mesh # if distributed.get_rank() == 0: # # only load model in rank0 to reduce network traffic and sync later # self.fsdp_checkpointer.load_model_during_init(model, is_ema=False) # if not hasattr(self, "fsdp_wrap_block_cls"): # raise ValueError("Networks does not have fsdp_wrap_block_cls attribute, please check the net definition") # fsdp_blocks_cls = self.fsdp_wrap_block_cls # fsdp_blocks_cls = ( # list(fsdp_blocks_cls) if isinstance(fsdp_blocks_cls, (list, tuple, set)) else [fsdp_blocks_cls] # ) # log.critical(f"Using FSDP blocks {fsdp_blocks_cls}") # log.critical(f"Using wrap policy {self.config.fsdp['policy']}") # if self.config.fsdp["policy"] == "size": # # Size based policy won't work for transformers because the tokenizers need to be accessible at multiple # # layers (input / output). This is handled by this sharding strategy. # min_num_params = self.config.fsdp["min_num_params"] # log.critical(f"Using {min_num_params} as the minimum number of parameters for auto-wrap policy") # log.info("If using a Transformer model. Please use the transformer wrap policy.") # wrap_policy = functools.partial(size_based_auto_wrap_policy, min_num_params=min_num_params) # else: # # Use the auto wrap policy for transformers # wrap_policy = functools.partial( # transformer_auto_wrap_policy, # transformer_layer_cls=set(fsdp_blocks_cls), # ) # tensor_kwargs = {"device": "cuda", "dtype": model.precision} # # Wrap the model with FSDP and attach it back to this class # self.model = FSDP( # model.to(**tensor_kwargs), # sync_module_states=True, # it can reduce network traffic by only loading model in rank0 and sync # sharding_strategy=strategy, # auto_wrap_policy=wrap_policy, # device_id=torch.cuda.current_device(), # device_mesh=device_mesh, # limit_all_gathers=True, # use_orig_params=True, # Do not flatten the parameter structure. Useful for layer_dependent lrs, etc. # ) # if self.config.act_ckpt_enabled: # # Apply activation checkpointing # apply_fsdp_checkpointing(self.model, list_block_cls=fsdp_blocks_cls) # # Clean up memory # torch.cuda.empty_cache() # def state_dict(self) -> Dict: # raise NotImplementedError("FSDPLlama does not support state_dict, use state_dict_model and FSDPCheckpointer") # @misc.timer("FSDP state_dict_model") # def state_dict_model(self) -> Dict: # """ # Get the model state_dict for checkpoint saving in the FSDP mode. # """ # with FSDP.summon_full_params(self.model): # pass # with FSDP.state_dict_type( # self.model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True) # ): # model_state = self.model.state_dict() # # No support for EMA yet. # ema_model_state = None # return { # "model": model_state, # "ema": ema_model_state, # } # def load_state_dict(self, state_dict: Dict, strict: bool = True, assign: bool = False) -> None: # raise NotImplementedError("FSDPLlama does not support load_state_dict, using FSDPCheckpointer") # def init_optimizer_scheduler( # self, optimizer_config: LazyDict, scheduler_config: LazyDict # ) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]: # """ # Initialize the optimizer and scheduler for FSDP model. # Args: # optimizer_config (LazyDict): The optimizer configuration. # scheduler_config (LazyDict): The scheduler configuration. # Returns: # tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]: The optimizer and scheduler. # """ # optimizer, scheduler = super().init_optimizer_scheduler(optimizer_config, scheduler_config) # self.fsdp_checkpointer.load_optim_scheduler_during_init( # self.model, # optimizer, # scheduler, # ) # return optimizer, scheduler # def get_ckpt_postfix(self) -> Tuple[str, int]: # """Get the checkpoint file postfix. check FSDPCheckpointer for more details # Returns: # postfix (str): The postfix of the checkpoint file. # replicate_idx, shard_idx (int), current gpu replicate_idx, shard_idx in FSDP \ # we will not save each ema model in each GPU, \ # ema model with same rate will be saved once # total_ema_num (int) # """ # replicate_idx, shard_idx = parallel_state.fsdp_device_mesh.get_coordinate() # # !!! EMA is not supported # if replicate_idx == 0: # return "", 0, shard_idx, 0 # return "", replicate_idx, shard_idx, 0