# 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. from typing import Any, Dict, Optional import torch import torch.nn as nn import transformer_engine as te from megatron.core import InferenceParams, ModelParallelConfig, parallel_state from megatron.core.tensor_parallel import gather_from_tensor_model_parallel_region from torch.distributed import ProcessGroup from torch.distributed import _functional_collectives as funcol from torch.distributed import broadcast, get_process_group_ranks from torch.nn.modules.module import _IncompatibleKeys from transformer_engine.pytorch.module.linear import Linear as LinearTE from transformer_engine.pytorch.module.rmsnorm import RMSNorm as RMSNormTE from cosmos_predict1.utils import log _ACTION_DIM = 8 from cosmos_predict1.autoregressive.modules.embedding import ( RotaryPositionEmbeddingPytorch, RotaryPositionEmbeddingPytorchV2, RotaryPositionEmbeddingTE, SinCosPosEmbAxisTE, get_pos_emb_on_this_cp_rank, get_pos_emb_on_this_sptp_rank, ) from cosmos_predict1.autoregressive.modules.linear import ColumnParallelLinear, TrainingVocabParallelEmbedding from cosmos_predict1.autoregressive.modules.mlp import TrainingMLP, compute_llama3_ffn_hidden_dim from cosmos_predict1.autoregressive.modules.normalization import create_norm from cosmos_predict1.autoregressive.training.modules.attention import ( GQA, create_group_causal_attn_mask, enable_different_context_dim_in_te_ca, enable_qk_normalization_in_te_mha, ) from cosmos_predict1.autoregressive.utils.checkpoint import process_state_dict, substrings_to_ignore from cosmos_predict1.autoregressive.utils.misc import maybe_convert_to_namespace from cosmos_predict1.autoregressive.utils.parallel import ( AllReduceBWDRMSNormTE, allreduce_layernorm_grads, sync_1d_parameters, ) _MLP_HIDDEN_DIM_DIVISOR = ( 4 # hidden dim of the action embedding layer is action_embedding_dim // _MLP_HIDDEN_DIM_DIVISOR ) _T5_NUM_TOKENS = 512 class TransformerBlock(nn.Module): """ A single transformer block consisting of an attention layer and a feed-forward layer. """ def __init__(self, layer_id: int, model_parallel: Optional[ModelParallelConfig] = None, args=None): """ Initializes the TransformerBlock module. Args: layer_id: The ID of the transformer block. args: The model arguments containing hyperparameters. """ super().__init__() args = maybe_convert_to_namespace(args) attention_args = { "n_heads": args["n_heads"], "n_kv_heads": args["n_kv_heads"], "dim": args["dim"], "context_dim": None, "max_batch_size": args["max_batch_size"], "max_seq_len": args["max_seq_len"], "inference": args["inference"], "flash_attn": args["flash_attn"], "use_qk_normalization": args["use_qk_normalization"], "attention_dropout": getattr(args, "attention_dropout", 0.0), "set_parallel_mode": args["set_parallel_mode"], "model_parallel": model_parallel, "attention_tp": args["attention_tp"], "causal_mask": args["causal_mask"], "head_dim": args["head_dim"], "fuse_qkv": getattr(args, "fuse_qkv", False), "precision": getattr(args, "precision", "bfloat16"), "attention_type": getattr(args, "attention_type", "self"), } self.attention = GQA(**attention_args) self.has_cross_attention = False self.cross_attention, self.cross_attention_norm = None, None if args["insert_cross_attn"] and layer_id % args["insert_cross_attn_every_k_layers"] == 0: self.has_cross_attention = True cross_attention_args = attention_args.copy() cross_attention_args.update( {"context_dim": args["context_dim"], "fuse_qkv": False, "attention_type": "cross"} ) self.cross_attention = GQA(**cross_attention_args) self.cross_attention_norm = create_norm(args["norm_type"], dim=args["dim"], eps=args["norm_eps"]) self.feed_forward = TrainingMLP( dim=args["dim"], hidden_dim=( compute_llama3_ffn_hidden_dim( dim=args["dim"], multiple_of=args["multiple_of"], ffn_dim_multiplier=args["ffn_dim_multiplier"] ) if args["ffn_hidden_size"] is None else args["ffn_hidden_size"] ), hidden_dropout=getattr(args, "hidden_dropout", 0.0), set_parallel_mode=args["set_parallel_mode"], model_parallel=model_parallel, inference=args["inference"], ) self.layer_id = layer_id self.num_layers = args["n_layers"] self.attention_norm = create_norm(args["norm_type"], dim=args["dim"], eps=args["norm_eps"]) self.ffn_norm = create_norm(args["norm_type"], dim=args["dim"], eps=args["norm_eps"]) # If `True`, then each transformer block init uses its layer ID, and if `False`, each uses the # total number of transformer blocks. Default is `True` (following the TorchTitan implementation of Llama3). if getattr(args, "depth_init", True): self.weight_init_std = 0.02 / (2 * (self.layer_id + 1)) ** 0.5 else: self.weight_init_std = 0.02 / (2 * self.num_layers) ** 0.5 def forward( self, x: torch.Tensor, rope: RotaryPositionEmbeddingPytorch, input_pos: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Performs the forward pass of the TransformerBlock module. Args: x: The input tensor. input_pos: The position of the current sequence. Used in inference (with KV cache) only. freqs_cis: The precomputed frequency values for rotary position embeddings. mask: The attention mask tensor. context (Optional[torch.Tensor]): The context tensor added via cross-attn. context_mask (Optional[torch.Tensor]): The context cross-attn mask tensor. Returns: The output tensor after applying the transformer block. """ # Apply attention and residual connection h = x + self.attention(self.attention_norm(x), rope=rope, input_pos=input_pos, mask=mask) # If insert cross-attention, apply CA and residual connection if self.has_cross_attention: h = h + self.cross_attention( self.cross_attention_norm(h), rope=rope, input_pos=input_pos, mask=context_mask, context=context ) # Apply feed-forward network and residual connection out = h + self.feed_forward(self.ffn_norm(h)) return out def init_weights(self): """ Initializes the weights of the transformer block. """ for norm in (self.attention_norm, self.ffn_norm): norm.reset_parameters() self.attention.init_weights(self.weight_init_std) self.feed_forward.init_weights(self.weight_init_std) if self.has_cross_attention: self.cross_attention_norm.reset_parameters() self.cross_attention.init_weights(self.weight_init_std) # zero-init the final output layer of cross-attention # nn.init.zeros_(self.cross_attention.wo.weight) class TransformerBlockTE(te.pytorch.TransformerLayer): """ Wrapper class over TE's `TransformerLayer`. Args: layer_id (int): The ID of the transformer block. args: The model arguments containing hyperparameters. """ def __init__( self, layer_id: int, args, tp_group: Optional[ProcessGroup] = None, set_parallel_mode: bool = False, attn_input_format: str = "bshd", ): attention_args = { "hidden_size": args["dim"], "ffn_hidden_size": ( compute_llama3_ffn_hidden_dim( dim=args["dim"], multiple_of=args["multiple_of"], ffn_dim_multiplier=args["ffn_dim_multiplier"] ) if args["ffn_hidden_size"] is None else args["ffn_hidden_size"] ), "num_attention_heads": args["n_heads"], "bias": False, "layernorm_epsilon": args["norm_eps"], "hidden_dropout": getattr(args, "hidden_dropout", 0.0), "attention_dropout": getattr(args, "attention_dropout", 0.0), "normalization": "RMSNorm", "activation": "swiglu", "attn_input_format": attn_input_format, "num_gqa_groups": args["n_kv_heads"], "fuse_wgrad_accumulation": False, "fuse_qkv_params": False, "tp_group": tp_group, "sequence_parallel": args["sequence_parallel"], "set_parallel_mode": set_parallel_mode, "layer_number": layer_id + 1, "self_attn_mask_type": "causal" if args["causal_mask"] else "no_mask", "kv_channels": args["head_dim"], # If None, te.pytorch.TransformerLayer defaults it to dim // n_heads "layer_type": "encoder", } self.has_cross_attention = False if args["insert_cross_attn"] and layer_id % args["insert_cross_attn_every_k_layers"] == 0: self.has_cross_attention = True attention_args["layer_type"] = "decoder" super().__init__(**attention_args) if args["use_qk_normalization"]: # Add QK normalization layers and replace the forward function of original Multi-Head Attention module with # our custom one to add QK normalization operations. enable_qk_normalization_in_te_mha(self.self_attention, norm_eps=args["norm_eps"], is_self_attn=True) if self.has_cross_attention: enable_qk_normalization_in_te_mha(self.inter_attention, norm_eps=args["norm_eps"], is_self_attn=False) if self.has_cross_attention: enable_different_context_dim_in_te_ca( self.inter_attention, context_dim=args["context_dim"], args=attention_args ) self.layer_id = layer_id self.num_layers = args["n_layers"] # If `True`, then each transformer block init uses its layer ID, and if `False`, each uses the # total number of transformer blocks. Default is `True` (following the TorchTitan implementation of Llama3). if getattr(args, "depth_init", True): self.weight_init_std = 0.02 / (2 * (self.layer_id + 1)) ** 0.5 else: self.weight_init_std = 0.02 / (2 * self.num_layers) ** 0.5 self.args = args self.inference = args["inference"] def set_inference_flag(self, flag: bool): """ Set the inference flag for the transformer layers. """ self.inference = flag def forward( self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, mask: Optional[torch.Tensor], inference_params: Optional[InferenceParams] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Custom forward to make sure we only pass relevant arguments to the forward pass of the `TransformerLayer`. Args: x (torch.Tensor): The input tensor. mask (Optional[torch.Tensor]): The attention mask tensor. inference_params (Optional[InferenceParams]): Inference parameters used for caching key-value pairs in the TE backend. It is not applicable for the PyTorch backend and should be set to None in that case. context (Optional[torch.Tensor]): The context tensor added via cross-attn. context_mask (Optional[torch.Tensor]): The context cross-attn mask tensor. Returns: torch.Tensor: The output tensor after applying the transformer block """ inference_params = None if not self.inference else inference_params output = super().forward( x, attention_mask=mask, rotary_pos_emb=rotary_pos_emb.to(x.device), inference_params=inference_params, encoder_output=context, enc_dec_attn_mask=context_mask, ) return output def init_weights(self): """ Initializes the weights of the transformer block. """ # Self Attention attn_layer = self.self_attention.layernorm_qkv for linear_weight in [attn_layer.query_weight, attn_layer.key_weight, attn_layer.value_weight]: nn.init.trunc_normal_(linear_weight, mean=0.0, std=0.02) nn.init.trunc_normal_(self.self_attention.proj.weight, mean=0.0, std=self.weight_init_std) # Cross Attention if self.has_cross_attention: nn.init.trunc_normal_(self.inter_attention.layernorm_query.query_weight, mean=0.0, std=0.02) nn.init.trunc_normal_(self.inter_attention.key_value.key_weight, mean=0.0, std=0.02) nn.init.trunc_normal_(self.inter_attention.key_value.value_weight, mean=0.0, std=0.02) # zero-init the final output layer of cross-attention if self.args["zero_init_cross_attn_proj"]: nn.init.zeros_(self.inter_attention.proj.weight) else: nn.init.trunc_normal_(self.inter_attention.proj.weight, mean=0.0, std=self.weight_init_std) # RMS Normalization for norm_weight in (self.layernorm_mlp.layer_norm_weight, self.self_attention.layernorm_qkv.layer_norm_weight): torch.nn.init.ones_(norm_weight) # In the case of QK Normalization, we also reset the parameters of the QK normalization layers. if self.args["use_qk_normalization"]: for norm_weight in [self.self_attention.q_norm.weight, self.self_attention.k_norm.weight]: torch.nn.init.ones_(norm_weight) # MLP for linear_weight in (self.layernorm_mlp.fc1_weight, self.layernorm_mlp.fc2_weight): nn.init.trunc_normal_(linear_weight, mean=0.0, std=self.weight_init_std) # The fc1_weight is a fused weight of w1 and w2 in the MLP of the PyTorch backend, where w1 is initialized with # a different std (0.02 by TorchTitan). So we re-initialize the w1 part of the fused weight below. split_point = self.layernorm_mlp.fc1_weight.shape[0] // 2 nn.init.trunc_normal_(self.layernorm_mlp.fc1_weight[:split_point], mean=0.0, std=0.02) class Transformer(nn.Module): """ The Transformer network consisting of transformer blocks. """ def __init__(self, params, model_parallel=None, tokenizer_config=None, init_weights: bool = True): """ Initializes the Transformer module. Args: params: The model parameters containing hyperparameters. model_parallel: The model parallel configuration. tokenizer_config: The model tokenizer configuration. init_weights (bool): Whether to initialize the weights of the transformer following TorchTitan's Llama3 initialization scheme. """ super().__init__() # Check if self.params is an OmegaConf DictConfig instance self.params = maybe_convert_to_namespace(params) self.vocab_size = params["vocab_size"] self.n_layers = params["n_layers"] self.precision = getattr(torch, params["precision"]) self.inference = params["inference"] self.backend = params["backend"] self.tokenizer_config = tokenizer_config self.model_parallel = model_parallel self.num_video_frames = params["num_video_frames"] self.token_emb_dropout = nn.Dropout(getattr(params, "embedding_dropout", 0.0)) tp_group = self._get_tp_group() # Sequence parallelism requires the first dimension to be the sequence dimension. When sequence parallelism # is enabled, we transpose the first two dimensions of the input tensor, and specify the format as "sbhd", # (sequence, batch, head, dim). Otherwise, the input format is "bshd" (batch, sequence, head, dim). self.attn_input_format = "bshd" if not params["sequence_parallel"] else "sbhd" # Token embeddings self.tok_embeddings = self._create_token_embeddings(self.model_parallel) self.rope_config = self._create_rope_config() if self.backend == "pytorch": self._initialize_pytorch_backend(model_parallel) elif self.backend == "transformer_engine": self._initialize_transformer_engine_backend(tp_group) else: raise ValueError(f"Unknown backend: {self.backend}") self.output = self._create_output_projection(model_parallel) # Action conditioning self.use_action_condition = getattr(params, "use_action_condition", False) if self.use_action_condition: self.action_dim = getattr( params, "action_dim", _ACTION_DIM ) # e.g., [Δx, Δy, Δz, rx, ry, rz, gripper_open, zero_pad] self.action_embedding_dim = self.params["action_embedding_dim"] # 1024 self.action_embedding_mode = getattr(params, "action_embedding_mode", "mlp") # Default to mlp mode self.group_causal_mask_mode = getattr( params, "group_causal_mask_mode", None ) # Default to None, 'causal' or 'group_diagonal' self.action_embedding_layers = self._create_action_projection() if params["sequence_parallel"]: if model_parallel is None: setattr(params, "sequence_parallel", False) log.critical("model_parallel is None. Disabling sequence parallelism.") self.sequence_parallel_enabled = False else: assert self.backend == "transformer_engine", f"Invalid backend: {self.backend} for sequence parallelism" assert ( params["tensor_model_parallel_size"] > 1 ), f"Invalid tensor_model_parallel_size: {params['tensor_model_parallel_size']}" self.sequence_parallel_enabled = True else: self.sequence_parallel_enabled = False if init_weights: self.init_weights() # Set default value for peft_last_n_layers and peft_every_n_layers self.peft_last_n_layers = getattr(params, "peft_last_n_layers", 0) self.peft_every_n_layers = getattr(params, "peft_every_n_layers", 0) if self.peft_last_n_layers > 0 or self.peft_every_n_layers > 0: self._setup_peft() # Freeze network parameters for finetuning w/ cross-attention self.has_cross_attention = getattr(params, "insert_cross_attn", False) if self.has_cross_attention: self.ca_every_k_layers = getattr(params, "insert_cross_attn_every_k_layers", 1) self.finetune_layers_with_cross_attn = getattr(params, "finetune_layers_with_cross_attn", False) self.finetune_layers_without_cross_attn = getattr(params, "finetune_layers_without_cross_attn", False) self._setup_cross_attn_ft() if self.params["apply_abs_pos_emb"]: self.pos_emb_config = self._create_abs_pos_emb_config() self.pos_emb, self.abs_pos_emb = self._initialize_abs_pos_emb() if self.attn_input_format == "sbhd": self.abs_pos_emb = self.abs_pos_emb.transpose(0, 1).contiguous() self._broadcast_pos_emb(self.abs_pos_emb, tp_group) def _initialize_pytorch_backend(self, model_parallel): self.layers = nn.ModuleList( [ TransformerBlock(layer_id, model_parallel, self.params).to(self.precision) for layer_id in range(self.n_layers) ] ) self.norm = create_norm(self.params["norm_type"], dim=self.params["dim"], eps=self.params["norm_eps"]).to( self.precision ) pytorch_rope_version = getattr(self.params, "pytorch_rope_version", "v2") if pytorch_rope_version == "v1": self.rope = RotaryPositionEmbeddingPytorch(**self.rope_config) elif pytorch_rope_version == "v2": training_type = self.tokenizer_config.training_type if self.tokenizer_config is not None else None self.rope = RotaryPositionEmbeddingPytorchV2( seq_len=self.params["max_seq_len"], training_type=training_type, **self.rope_config ) self._broadcast_pos_emb(self.rope.cos_cached, tp_group=self._get_tp_group()) self._broadcast_pos_emb(self.rope.sin_cached, tp_group=self._get_tp_group()) else: raise ValueError(f"Unknown pytorch_rope_version: {pytorch_rope_version}") self.causal_mask = torch.tril( torch.ones(self.params["max_seq_len"], self.params["max_seq_len"], dtype=torch.bool) ).cuda() def _initialize_transformer_engine_backend(self, tp_group): self.layers = self._create_transformer_layers(tp_group) if self.params["sequence_parallel"]: tp_group = parallel_state.get_tensor_model_parallel_group() self.norm = AllReduceBWDRMSNormTE( self.params["dim"], process_group=tp_group, eps=self.params["norm_eps"], sequence_parallel=True, ).to(self.precision) else: self.norm = RMSNormTE(self.params["dim"], eps=self.params["norm_eps"]).to(self.precision) self.rope, self.rotary_pos_emb = self._initialize_rope() self._broadcast_pos_emb(self.rotary_pos_emb, tp_group) def _create_rope_config(self) -> Dict: shape_map = { "3D": self.params["video_latent_shape"], "2D": self.params["image_latent_shape"], "1D": None, } latent_shape = shape_map.get(self.params["rope_dim"], None) head_dim = self.params["head_dim"] if head_dim is None: head_dim = self.params["dim"] // self.params["n_heads"] return { "dim": head_dim, "max_position_embeddings": self.params["max_seq_len"], "original_max_position_embeddings": self.params["original_seq_len"], "rope_theta": self.params["rope_theta"], "apply_yarn": self.params["apply_yarn"], "scale": self.params["yarn_scale"], "beta_fast": self.params["yarn_beta_fast"], "beta_slow": self.params["yarn_beta_slow"], "rope_dim": self.params["rope_dim"], "latent_shape": latent_shape, "original_latent_shape": self.params["original_latent_shape"], "pad_to_multiple_of": self.params["pad_to_multiple_of"], } def _create_abs_pos_emb_config(self): shape_map = { "3D": self.params["video_latent_shape"], "2D": self.params["image_latent_shape"], "1D": None, } latent_shape = shape_map.get(self.params["rope_dim"], None) return { "dim": self.params["dim"], "latent_shape": latent_shape, "pad_to_multiple_of": self.params["pad_to_multiple_of"], } def _create_token_embeddings(self, model_parallel=None, vocab_size: int = None): """ Create token embeddings. Args: model_parallel: The model parallel configuration. Returns: nn.Module: Token embeddings module. """ if vocab_size is None: vocab_size = self.params["vocab_size"] tp_size = self.params["tensor_model_parallel_size"] if tp_size > 1: # For inference in the PyTorch backend, we use PyTorch's allreduce (tracable) in the forward pass to enable torch.compile. use_inference_allreduce = self.inference and self.params["backend"] == "pytorch" emb = TrainingVocabParallelEmbedding( vocab_size, self.params["dim"], init_method=lambda x: x, config=model_parallel, sequence_parallel=self.params["sequence_parallel"], batch_first=not self.params["sequence_parallel"], use_inference_allreduce=use_inference_allreduce, ).to(self.precision) return emb else: return nn.Embedding(vocab_size, self.params["dim"]).to(self.precision) def _create_action_projection(self): """ Create the action projection layer. Returns: nn.Module: Action projection layer. """ assert self.action_embedding_mode == "mlp", f"Invalid action embedding mode: {self.action_embedding_mode}" # This method is not working well. (option 1. default) exp102e hidden_dim = self.action_embedding_dim // _MLP_HIDDEN_DIM_DIVISOR action_embedding_layers = nn.Sequential( nn.Linear(self.action_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, self.action_embedding_dim), ) return action_embedding_layers def _get_tp_group( self, ): """ Get tensor parallel process group if applicable. Returns: torch.distributed.ProcessGroup or None: Tensor parallel process group if tensor parallelism is enabled, else None. """ if self.params["tensor_model_parallel_size"] > 1: tp_group = parallel_state.get_tensor_model_parallel_group() log.info(f"Using tensor model parallel group: {tp_group}") return tp_group return None def _create_transformer_layers(self, tp_group): """ Create the transformer layers. Args: tp_group (torch.distributed.ProcessGroup or None): Tensor parallel process group. Returns: nn.ModuleList: List of transformer layers. """ return nn.ModuleList( [ TransformerBlockTE( layer_id, self.params, tp_group, set_parallel_mode=self.params["set_parallel_mode"], attn_input_format=self.attn_input_format, ).to(self.precision) for layer_id in range(self.params["n_layers"]) ] ) def _create_output_projection(self, model_parallel=None, vocab_size: int = None): """ Create the output projection layer. Args: model_parallel: The model parallel configuration. vocab_size (int): Vocabulary size (to override the default vocab size). Returns: LinearTE: Output projection layer. """ if vocab_size is None: vocab_size = self.params["vocab_size"] if self.params["tensor_model_parallel_size"] > 1: if self.params["backend"] == "pytorch" and self.inference: tp_size = self.params["tensor_model_parallel_size"] layer = nn.Linear(self.params["dim"], vocab_size // tp_size, bias=False).to(self.precision) return layer else: layer = ColumnParallelLinear( self.params["dim"], vocab_size, bias=False, gather_output=False, init_method=lambda x: x, config=model_parallel, ).to(self.precision) return layer else: # No Tensor Parallelism if self.params["backend"] == "pytorch": return nn.Linear(self.params["dim"], vocab_size, bias=False).to(self.precision) elif self.params["backend"] == "transformer_engine": return LinearTE(self.params["dim"], vocab_size, bias=False).to(self.precision) else: raise ValueError("Unknown backend: " + self.params["backend"]) def _initialize_rope( self, ): """ Initialize the rotary position embedding. Returns: tuple: (RotaryPositionEmbeddingTE, torch.Tensor) The RoPE module and the rotary position embeddings. """ rope = RotaryPositionEmbeddingTE(**self.rope_config) training_type = self.tokenizer_config.training_type if self.tokenizer_config is not None else None rotary_pos_emb = rope.forward(seq_len=self.params["max_seq_len"], training_type=training_type) return rope, rotary_pos_emb def _initialize_abs_pos_emb(self): pos_emb = SinCosPosEmbAxisTE(**self.pos_emb_config) training_type = self.tokenizer_config.training_type if self.tokenizer_config is not None else None abs_pos_emb = pos_emb.forward(training_type=training_type) return pos_emb, abs_pos_emb def _broadcast_pos_emb(self, pos_emb, tp_group): """ Broadcast the position embeddings across the tensor parallel group. Args: pos_emb (torch.Tensor): Position embeddings to broadcast. tp_group (torch.distributed.ProcessGroup or None): Tensor parallel process group. """ if self.params["tensor_model_parallel_size"] > 1: broadcast(pos_emb, min(get_process_group_ranks(tp_group)), group=tp_group) def _setup_peft(self): """ Set up Parameter Efficient Fine-Tuning (PEFT) by selectively freezing and unfreezing layers. This method configures the model for fine-tuning by: 1. Freezing all parameters in the model. 2. Unfreezing the embedding, normalization and output layers. 3. Unfreezing the first and last (peft_last_n_layers - 1) transformer layers if peft_last_n_layers is set, or unfreezing every n layers (flamingo style) if peft_every_n_layers is set. """ # Ensure only one of peft_last_n_layers and peft_every_n_layers is set assert ( self.peft_last_n_layers == 0 or self.peft_every_n_layers == 0 ), "Only one of peft_last_n_layers and peft_every_n_layers can be set." # First, freeze all parameters for param in self.parameters(): param.requires_grad = False # Unfreeze embedding, normalization and output layers for param in self.tok_embeddings.parameters(): param.requires_grad = True for param in self.norm.parameters(): param.requires_grad = True for param in self.output.parameters(): param.requires_grad = True # PEFT last n layers if self.peft_last_n_layers > 0: # Ensure peft_last_n_layers is at least 2 assert self.peft_last_n_layers >= 2, "peft_last_n_layers must be at least 2" # Unfreeze specific transformer layers total_layers = len(self.layers) for i, layer in enumerate(self.layers): if i == 0 or i >= total_layers - self.peft_last_n_layers + 1: # Unfreeze the first layer and the last (peft_last_n_layers - 1) layers for param in layer.parameters(): param.requires_grad = True log.info( f"PEFT setup complete. Trainable components: embeddings, un-embedding, normalization layer, " f"first transformer layer, last {self.peft_last_n_layers - 1} transformer layers." ) # PEFT every n layers (flamingo style, e.g. every 4 layers = layer 0,1,2,4,5,6,... frozen, layer 3,7,11,... is trainable) else: trainable_layers = [] for i, layer in enumerate(self.layers, 1): if i % self.peft_every_n_layers == 0: for param in layer.parameters(): param.requires_grad = True trainable_layers.append(i - 1) log.info( f"PEFT setup complete. Trainable components: embeddings, un-embedding, normalization layer, " f"every {self.peft_every_n_layers} transformer layers (layer idx {trainable_layers}; total {len(trainable_layers)} layers)." ) def _setup_cross_attn_ft(self): """ Set up Cross Attention Fine-Tuning by selectively freezing and unfreezing layers. This method configures the model for fine-tuning by: 1. Freezing all parameters in the model. 2. Unfreezing the embedding, normalization and output layers. 3. Unfreezing all the added cross-attention layers. 4. If `finetune_layers_with_cross_attn` is True, unfreeze the transformer layers for layers with cross attention. 5. If `finetune_layers_without_cross_attn` is True, unfreeze the transformer layers for layers without cross attention. 6. If 'use_action_condition' is True, unfreeze the action embedding layers. """ assert self.has_cross_attention, "Must insert cross-attention layers for finetuning." finetune_layer_num = 0 # First, freeze all parameters for param in self.parameters(): param.requires_grad = False # Unfreeze embedding, normalization and output layers for param in self.tok_embeddings.parameters(): param.requires_grad = True for param in self.norm.parameters(): param.requires_grad = True for param in self.output.parameters(): param.requires_grad = True # Unfreeze all the added cross-attention layers total_layers = len(self.layers) for i, layer in enumerate(self.layers): if i % self.ca_every_k_layers == 0: if self.params["backend"] == "pytorch": for param in layer.cross_attention.parameters(): param.requires_grad = True elif self.params["backend"] == "transformer_engine": for param in layer.inter_attention.parameters(): param.requires_grad = True else: raise ValueError("Unknown backend: " + self.params["backend"]) # Unfreeze the transformer layers for layers with cross attention if self.finetune_layers_with_cross_attn: for i, layer in enumerate(self.layers): if i % self.ca_every_k_layers == 0: for param in layer.parameters(): param.requires_grad = True finetune_layer_num += 1 # Unfreeze the transformer layers for layers without cross attention if self.finetune_layers_without_cross_attn: for i, layer in enumerate(self.layers): if i % self.ca_every_k_layers != 0: for param in layer.parameters(): param.requires_grad = True finetune_layer_num += 1 # Unfreeze the action embedding layers if self.use_action_condition: for param in self.action_embedding_layers.parameters(): param.requires_grad = True log.info( f"cross attention finetune setup complete. Trainable components: cross-attention layer, " f"fully trainable transformer layer number is {finetune_layer_num}." ) def enable_context_parallel(self, cp_group: ProcessGroup): """ Enable context parallelism for the transformer model. This method sets up context parallelism by configuring the context parallel group and updating each transformer layer to support context parallelism. Args: cp_group (ProcessGroup): The process group for context parallelism. Notes: - Updates the model's context parallel group and size. - Configures each transformer layer for context parallelism. - Enables context parallelism for the rotary position embedding if using the transformer engine backend. """ cp_ranks = get_process_group_ranks(cp_group) cp_size = len(cp_ranks) # Set these attributes for spliting the data after embedding. self.cp_group = cp_group # Set these attributes for computing the loss. self.cp_size = cp_size for layer_idx, layer in enumerate(self.layers): if isinstance(layer, TransformerBlockTE): layer.set_context_parallel_group(cp_group, cp_ranks, torch.cuda.Stream()) elif hasattr(layer, "module") and isinstance(layer.module, TransformerBlockTE): layer.module.set_context_parallel_group(cp_group, cp_ranks, torch.cuda.Stream()) else: log.warning(f"Layer {layer_idx} does not support context parallelism") def set_inference_flag(self, flag: bool): """ Set the inference flag for the transformer layers. """ log.info(f"Setting inference flag to {flag}") self.inference = flag if self.inference: self.eval() if self.params["backend"] == "pytorch": for layer in self.layers: layer.attention.set_inference_flag(flag) elif self.params["backend"] == "transformer_engine": for layer in self.layers: layer.set_inference_flag(flag) self._maybe_change_sequence_parallel_status(enable=False) def _maybe_change_sequence_parallel_status(self, enable: bool): """ Change the sequence parallel status of the transformer layers. """ if enable and not self.sequence_parallel_enabled: for name, module in self.named_modules(): if hasattr(module, "sequence_parallel"): assert isinstance( module.sequence_parallel, bool ), f"Invalid type of {name}: {type(module.sequence_parallel)}" setattr(module, "sequence_parallel", True) self.sequence_parallel_enabled = True elif not enable and self.sequence_parallel_enabled: for name, module in self.named_modules(): if hasattr(module, "sequence_parallel"): assert isinstance( module.sequence_parallel, bool ), f"Invalid type of {name}: {type(module.sequence_parallel)}" setattr(module, "sequence_parallel", False) self.sequence_parallel_enabled = False def forward( self, tokens: Optional[torch.Tensor] = None, input_pos: Optional[torch.Tensor] = None, inference_params: Optional[InferenceParams] = None, token_embeddings: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, action: Optional[torch.Tensor] = None, total_seq_len: Optional[int] = None, return_hidden_states: bool = False, mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Performs the forward pass of the Transformer module. Args: tokens (torch.Tensor, optional): The input tensor of token IDs. input_pos (Optional[torch.Tensor]): The position of the current sequence. Used in inference with KV cache. PyTorch backend only. inference_params (InferenceParams, optional): Parameters for inference. token_embeddings (torch.Tensor, optional): Precomputed token embeddings. If provided, tokens should be None. context (Optional[torch.Tensor]): The context tensor added via cross-attn. context_mask (Optional[torch.Tensor]): The context cross-attn mask tensor. action (Optional[torch.Tensor]): The robot action tensor for conditioning. total_seq_len (Optional[int]): The total sequence length (before applying context parallelism). return_hidden_states (bool): Whether to return hidden states. Returns: The output tensor after applying the transformer layers. """ # Turn on/off sequence parallelism based on the training status self._maybe_change_sequence_parallel_status(enable=self.training and self.params["sequence_parallel"]) # Token embeddings assert ( tokens is None or token_embeddings is None ), "Either tokens or token_embeddings should be provided, not both." if token_embeddings is None: seq_len = tokens.shape[1] h = self.token_emb_dropout(self.tok_embeddings(tokens)) else: seq_len = token_embeddings.shape[1] h = self.token_emb_dropout(token_embeddings) if mask is None: # Create attention mask mask = self._create_attention_mask(input_pos=input_pos) # Action embedding if self.use_action_condition and action is not None: assert self.action_embedding_mode == "mlp", f"Invalid action embedding mode: {self.action_embedding_mode}" # change action type to bfloat16, of shape [batch_size, action_dim] action = action.to(torch.bfloat16) # action_emb shape: [batch_size, action_dim, action_embedding_dim] action_emb = self.action_embedding_layers(action).unsqueeze(1).repeat(1, self.action_dim, 1) # Use action_emb as context if self.params["concat_action_to_context"]: context = torch.zeros( (action_emb.shape[0], _T5_NUM_TOKENS, self.action_embedding_dim), device=h.device, dtype=h.dtype ) # context[:, -1, :] = action_emb[:, 0, :] # overwrite the last token with action_emb context = torch.cat([context, action_emb[:, 0:1, :]], dim=1) else: context = action_emb # [batch_size, action_dim, action_embedding_dim] # Create context mask if self.group_causal_mask_mode is not None: num_temporal_groups = self.num_video_frames - 1 # number of latent frames num_query_per_group = seq_len // num_temporal_groups # number of latent tokens per frame num_key_per_group = self.action_dim // num_temporal_groups context_mask = create_group_causal_attn_mask( num_temporal_groups=num_temporal_groups, num_query_per_group=num_query_per_group, num_key_per_group=num_key_per_group, mode=self.group_causal_mask_mode, ) # [L (query), S (key)] context_mask = context_mask.unsqueeze(0) # [1, L (query), S (key)] context_mask = context_mask.repeat(context.shape[0], 1, 1) # [batch_size, L (query), S (key)] context_mask = context_mask.to(context.device) else: context_mask = torch.ones( (context.shape[0], context.shape[1]), device=context.device, dtype=torch.bool ) # [batch_size, action_dim] # Prepare layer arguments layer_kwargs = self._prepare_layer_kwargs( total_seq_len=total_seq_len, input_pos=input_pos, mask=mask, inference_params=inference_params, context=context, context_mask=context_mask, ) # Apply transformer layers for layer in self.layers: if self.params["apply_abs_pos_emb"]: h = self.apply_abs_pos_emb(h, input_pos=input_pos, total_seq_len=total_seq_len) h = layer(h, **layer_kwargs) # Apply final layer normalization h = self.norm(h) if return_hidden_states: return h # Output linear projection output = self.output(h) output = self.process_output(output) return output def process_output(self, output: torch.Tensor) -> torch.Tensor: """ Adjusts the shape and layout of tensor based on tensor parallelism and attention input format. The function performs two operations: 1. If the tensor model parallelism is enabled (`tensor_model_parallel_size > 1`), it gathers the tensor from the tensor-parallel regions and reshapes it accordingly. 2. If the attention input format is `"sbhd"` (Sequence, Batch, Hidden Dimension), it transposes the tensor to the format `(Batch, Sequence, Hidden Dimension)` for further processing. Args: output [torch.Tensor]: The tensor before modification. Returns: output [torch.Tensor]: The tensor after modification. """ if self.params["tensor_model_parallel_size"] > 1: if self.params["backend"] == "pytorch" and self.inference: # Use PyTorch all gather output = funcol.all_gather_tensor( output, gather_dim=-1, group=parallel_state.get_tensor_model_parallel_group() ) else: # [*, *, hidden_dim // tp_size] --> [*, *, hidden_dim] output = gather_from_tensor_model_parallel_region(output) if self.attn_input_format == "sbhd": # [seq_len, batch_size, hidden_dim] --> [batch_size, seq_len, hidden_dim] output = output.transpose(0, 1).contiguous() return output def _create_attention_mask(self, input_pos: Optional[torch.Tensor]) -> Optional[torch.Tensor]: """ Creates an attention mask for the transformer layers. Args: input_pos[torch.Tensor]: The position of input sequence (used for inference only). Returns: Optional[torch.Tensor]: The attention mask, or None for causal mask. """ if self.backend == "pytorch" and self.inference: assert input_pos is not None, "input_pos must be provided for inference" mask = self.causal_mask[input_pos] return mask else: return None # None means causal mask def _prepare_layer_kwargs( self, total_seq_len: Optional[int], input_pos: Optional[torch.Tensor], mask: Optional[torch.Tensor], inference_params: Optional[InferenceParams], context: Optional[torch.Tensor], context_mask: Optional[torch.Tensor], ) -> Dict[str, Any]: """ Prepares the keyword arguments for transformer layers. Args: total_seq_len (Optional[int]): The total sequence length (before applying context parallelism). seq_len (Optional[int]): The length of the input sequence. input_pos (Optional[torch.Tensor]): The position of the current sequence. mask (Optional[torch.Tensor]): The attention mask. inference_params (Optional[InferenceParams]): Parameters for inference. context (Optional[torch.Tensor]): The context tensor added via cross-attn. context_mask (Optional[torch.Tensor]): The context cross-attn mask tensor. Returns: Dict[str, Any]: A dictionary of keyword arguments for the transformer layers. """ if context is not None: context = context.to(self.precision) if self.attn_input_format == "sbhd": context = context.transpose(0, 1).contiguous() if self.backend == "pytorch": if isinstance(mask, torch.Tensor) and mask.ndim == 2: mask = mask[None, None, :, :] if isinstance(context_mask, torch.Tensor) and context_mask.ndim == 2: context_mask = context_mask[None, None, :, :] layer_kwargs = { "mask": mask, "context": context, "context_mask": context_mask, } if self.backend == "pytorch": layer_kwargs["input_pos"] = input_pos layer_kwargs["rope"] = self.rope elif self.backend == "transformer_engine": rotary_pos_emb = self.rotary_pos_emb try: cp_size = parallel_state.get_context_parallel_world_size() except (AssertionError, RuntimeError): # Fallback if context parallel group isn't initialized cp_size = 1 log.warning("Context parallel group not initialized, falling back to size 1") else: cp_size = 1 if cp_size > 1: assert input_pos is None, "input_pos must be None for context parallelism" rotary_pos_emb = rotary_pos_emb[:total_seq_len] rotary_pos_emb = get_pos_emb_on_this_cp_rank(rotary_pos_emb, 0) layer_kwargs["rotary_pos_emb"] = rotary_pos_emb layer_kwargs["inference_params"] = inference_params return layer_kwargs def apply_abs_pos_emb( self, x: torch.Tensor, input_pos: int = None, total_seq_len: Optional[int] = None ) -> torch.Tensor: """ Applies the absolute position embeddings to the input tensor. """ abs_pos_emb = self.abs_pos_emb if total_seq_len is not None: # Truncate the absolute position embeddings to the total sequence length abs_pos_emb = ( abs_pos_emb[:total_seq_len, :, :] if self.attn_input_format == "sbhd" else abs_pos_emb[:, :total_seq_len, :] ) cp_size = parallel_state.get_context_parallel_world_size() if self.training else 1 if cp_size > 1: assert input_pos is None seq_dim = 0 if self.attn_input_format == "sbhd" else 1 abs_pos_emb = get_pos_emb_on_this_cp_rank(abs_pos_emb, seq_dim=seq_dim) if self.attn_input_format == "sbhd": if self.sequence_parallel_enabled: # Training assert input_pos is None, "input_pos must be None when training with sequence parallelism" abs_pos_emb = get_pos_emb_on_this_sptp_rank(abs_pos_emb, seq_dim=0) else: # Inference or Evaluation abs_pos_emb = abs_pos_emb[input_pos, :, :] if input_pos is not None else abs_pos_emb else: abs_pos_emb = abs_pos_emb[:, input_pos, :] if input_pos is not None else abs_pos_emb return x + abs_pos_emb @torch.no_grad() def expand_vocab( self, new_vocab_size: int, init_method: str = "gaussian", multiple_of=64, expand_output_layer=True ): """ Expands the vocabulary of the model to the new size. Args: new_vocab_size (int): The new vocabulary size. init_method (str): The initialization method for new embeddings. Can be "zero" or "gaussian". Default is "gaussian". multiple_of (int): The new vocabulary size must be a multiple of this value. Defaults to 64 to fully leverage the power of NVIDIA TensorCore (source 1: https://x.com/karpathy/status/1621578354024677377, source 2: https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc) expand_output_layer (bool): Whether to also expand the output layer. Defaults to True. Returns: None """ tp_size = self.params["tensor_model_parallel_size"] if new_vocab_size <= self.vocab_size: raise ValueError( f"New vocabulary size ({new_vocab_size}) must be " f"larger than current size ({self.vocab_size})" ) if new_vocab_size % multiple_of != 0: log.critical(f"New vocabulary size must be a multiple of {multiple_of}. Obtained {new_vocab_size}.") new_vocab_size = (new_vocab_size // multiple_of + 1) * multiple_of log.critical(f"Rounded vocabulary size to {new_vocab_size}.") # Resize token embeddings old_embeddings = self.tok_embeddings old_embeddings_requires_grad = old_embeddings.weight.requires_grad tensor_kwargs = {"device": old_embeddings.weight.device, "dtype": old_embeddings.weight.dtype} self.tok_embeddings = self._create_token_embeddings( model_parallel=self.model_parallel, vocab_size=new_vocab_size ).to(**tensor_kwargs) # Initialize new embeddings if init_method not in ["zero", "gaussian"]: raise ValueError(f"Unknown initialization method: {init_method}") # The default initialization of nn.Embedding is Gaussian, so we don't need to do anything # if init_method == "gaussian". Only if init_method == "zero", we need to zero out the new embeddings. if init_method == "zero": self.tok_embeddings.weight.data[self.vocab_size // tp_size :].zero_() # Copy old embeddings log.info( f"old_embeddings: {old_embeddings.weight.data.shape}, new_embeddings: {self.tok_embeddings.weight.data.shape}, vocab_size: {self.vocab_size}" ) self.tok_embeddings.weight.data[: self.vocab_size // tp_size] = old_embeddings.weight.data self.tok_embeddings.weight.requires_grad = old_embeddings_requires_grad # Resize output layer old_output = self.output old_output_requires_grad = old_output.weight.requires_grad self.output = self._create_output_projection( self.model_parallel, vocab_size=new_vocab_size if expand_output_layer else None ) # Initialize new output weights if init_method == "zero": self.output.weight.data[self.vocab_size // tp_size :].zero_() elif init_method == "gaussian": # Follows the parameter initialization in TorchTitan: # https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/llama/model.py final_out_std = self.params["dim"] ** -0.5 cutoff_factor = 3 nn.init.trunc_normal_( self.output.weight, mean=0.0, std=final_out_std, a=-cutoff_factor * final_out_std, b=cutoff_factor * final_out_std, ) # Copy old output weights self.output.weight.data[: self.vocab_size // tp_size] = old_output.weight.data self.output.weight.requires_grad = old_output_requires_grad # Update vocab size self.vocab_size = new_vocab_size log.critical(f"Expanded vocabulary size to {new_vocab_size}") def init_weights(self): """ [Note: On ``init_weights`` vs. ``reset_parameters`` (copied from github.com/pytorch/torchtitan)] Modules may define ``reset_parameters`` to initialize parameter values. ``reset_parameters`` is meant to only initialize directly owned parameters/buffers, not those of their child modules, and it can be used to give the initial values for these tensors. Separately, users may want custom initialization for their modules, different from that in ``reset_parameters``. For this, we define ``init_weights``. We only call it in the constructor of this ``Transformer`` root module to avoid reinitializing tensors. """ nn.init.normal_(self.tok_embeddings.weight) for layer in self.layers: layer.init_weights() if self.backend == "pytorch": self.norm.reset_parameters() elif self.backend == "transformer_engine": nn.init.ones_(self.norm.weight) else: raise ValueError(f"Unknown backend: {self.backend}") final_out_std = self.params["dim"] ** -0.5 cutoff_factor = 3 nn.init.trunc_normal_( self.output.weight, mean=0.0, std=final_out_std, a=-cutoff_factor * final_out_std, b=cutoff_factor * final_out_std, ) if self.use_action_condition: for layer in self.action_embedding_layers: if isinstance(layer, nn.Linear): nn.init.xavier_uniform_(layer.weight) nn.init.zeros_(layer.bias) 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) if strict: 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 len(actual_missing_keys) > 0 or len(unexpected_keys) > 0: raise ValueError(f"Missing keys: {actual_missing_keys}\n\nUnexpected keys: {unexpected_keys}") missing_keys = actual_missing_keys return _IncompatibleKeys(missing_keys, unexpected_keys) def on_after_backward(self, *args, **kwargs): """ All-reduce layernorm grads for tensor/sequence parallelism. Reference implementation: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/distributed/finalize_model_grads.py """ allreduce_layernorm_grads( [self], tensor_model_parallel_size=self.params["tensor_model_parallel_size"], sequence_parallel=self.params["sequence_parallel"], ) 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. """ if self.params["sync_1d_parameters"]: if self.params["tensor_model_parallel_size"] > 1: sync_1d_parameters(self, process_group=parallel_state.get_tensor_model_parallel_group()) if self.params["context_parallel_size"] > 1: sync_1d_parameters(self, process_group=parallel_state.get_context_parallel_group())