import math from dataclasses import dataclass, replace import torch from torch import nn from torch.distributed._tensor import DTensor, Replicate from torchtitan.protocols.train_spec import ModelProtocol from llm_training.layers.attention import DiffAttentionV2 from llm_training.layers.ffn_clip import ffn_clip_loss as _ffn_clip_loss_fn from llm_training.layers.mhc import MHCLayer from llm_training.layers.moe import FeedForward, MoE from llm_training.motif.model.args import MotifModelArgs, MotifMoEModelArgs from llm_training.motif.utils import shift_left_and_pad from llm_training.tools.kernels import activation from llm_training.tools.logging import logger def _ulysses_a2a_impl( x: torch.Tensor, group: "torch.distributed.ProcessGroup", scatter_dim: int, gather_dim: int, ) -> torch.Tensor: """Ulysses all-to-all: scatter along scatter_dim, gather along gather_dim.""" world_size = group.size() if world_size == 1: return x input_list = [t.contiguous() for t in x.chunk(world_size, dim=scatter_dim)] output_list = [torch.empty_like(t) for t in input_list] torch.distributed.all_to_all(output_list, input_list, group=group) return torch.cat(output_list, dim=gather_dim) class _UlyssesA2A(torch.autograd.Function): """Autograd-aware Ulysses all-to-all. torch.distributed.all_to_all is a non-autograd collective; calling it directly inside a forward leaves the backward graph severed at the boundary, so gradients on every weight upstream of the a2a are zero. Wrapping it as an autograd.Function makes the backward run a transposed a2a (swap scatter/gather dims) that correctly routes the upstream gradient back to the producing rank. """ @staticmethod def forward(ctx, x, group, scatter_dim, gather_dim): ctx.group = group ctx.scatter_dim = scatter_dim ctx.gather_dim = gather_dim return _ulysses_a2a_impl(x, group, scatter_dim, gather_dim) @staticmethod def backward(ctx, grad_output): # Inverse of (scatter on A, gather on B) is (scatter on B, gather on A). return ( _ulysses_a2a_impl(grad_output, ctx.group, ctx.gather_dim, ctx.scatter_dim), None, None, None, ) def _ulysses_a2a(x, group, scatter_dim, gather_dim): return _UlyssesA2A.apply(x, group, scatter_dim, gather_dim) def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor: """ Precompute the frequency tensor for complex exponentials (cis) with given dimensions. This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64 data type. Args: dim (int): Dimension of the frequency tensor. end (int): End index for precomputing frequencies. theta (float | None): Scaling factor for frequency computation. Defaults to 10000.0. Returns: torch.Tensor: Precomputed frequency tensor with complex exponentials. """ freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) t = torch.arange(end, device=freqs.device) freqs = torch.outer(t, freqs).float() freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 return freqs_cis def precompute_freqs_cis_yarn( dim: int, end: int, theta: float, original_seq_len: int, rope_factor: float, beta_fast: int, beta_slow: int, ) -> torch.Tensor: """YaRN-scaled RoPE frequencies (DeepSeek-V3 style). Extends context from ``original_seq_len`` to ``end`` by blending the wavelength-scaled and base frequency bands over the NTK "correction range" defined by (beta_fast, beta_slow). When ``end <= original_seq_len`` this reduces to plain RoPE. Adapted from torchtitan.models.deepseek_v3.model.model.precompute_freqs_cis. """ def find_correction_dim(num_rotations: float, d: int, base: float, max_seq_len: int) -> float: return d * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) def find_correction_range( low_rot: float, high_rot: float, d: int, base: float, max_seq_len: int ) -> tuple[int, int]: low = math.floor(find_correction_dim(low_rot, d, base, max_seq_len)) high = math.ceil(find_correction_dim(high_rot, d, base, max_seq_len)) return max(low, 0), min(high, d - 1) def linear_ramp_factor(lo: float, hi: float, d: int) -> torch.Tensor: if lo == hi: hi += 0.001 ramp = (torch.arange(d, dtype=torch.float32) - lo) / (hi - lo) return torch.clamp(ramp, 0, 1) freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) if end > original_seq_len: low, high = find_correction_range(beta_fast, beta_slow, dim, theta, original_seq_len) smooth = 1 - linear_ramp_factor(low, high, dim // 2) freqs = freqs / rope_factor * (1 - smooth) + freqs * smooth t = torch.arange(end) freqs = torch.outer(t, freqs).float() freqs_cis = torch.polar(torch.ones_like(freqs), freqs) return freqs_cis def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor: """ Reshape frequency tensor for broadcasting it with another tensor. This function reshapes the frequency tensor to have the same shape as the target tensor 'x' for the purpose of broadcasting the frequency tensor during element-wise operations. The input freqs_cis tensor is assumed to be of shape (max_seqlen, dim), and the first seqlen elements will be sliced, but dim must match x. Args: freqs_cis (torch.Tensor): Frequency tensor to be reshaped. x (torch.Tensor): Target tensor for broadcasting compatibility. Returns: torch.Tensor: Reshaped frequency tensor. """ ndim = x.ndim assert ndim > 1 seqlen = x.shape[1] freqs_cis = freqs_cis[0:seqlen] assert freqs_cis.shape == (seqlen, x.shape[-1]), f"freqs_cis.shape: {freqs_cis.shape}, x.shape: {x.shape}" shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] return freqs_cis.view(*shape) def apply_rotary_emb_single( x: torch.Tensor, freqs_cis: torch.Tensor, ) -> torch.Tensor: """ Apply rotary embeddings to a single tensor using the given frequency tensor. Args: x (torch.Tensor): Input tensor to apply rotary embeddings. freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials. Returns: torch.Tensor: Modified tensor with rotary embeddings. """ x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) freqs_cis = reshape_for_broadcast(freqs_cis, x_) x_out = torch.view_as_real(x_ * freqs_cis).flatten(3) return x_out.type_as(x) def _compute_window_size(layer_idx: int, model_args: MotifModelArgs) -> tuple[int, int]: """Return the flash-attention window_size tuple for this layer. full-first interleave pattern (sliding_window_period=P): layer 0, P, 2P, ... → full (-1, -1) all others → SWA (w, 0) Right window is 0 (no future tokens) because attention is causal. """ sliding_window = getattr(model_args, "sliding_window", None) if sliding_window is None: return (-1, -1) pattern = getattr(model_args, "sliding_window_pattern", "interleave") period = getattr(model_args, "sliding_window_period", 2) if pattern == "all": return (sliding_window, 0) if pattern == "interleave" and layer_idx % period != 0: return (sliding_window, 0) return (-1, -1) class GDLAttention(nn.Module): """Grouped Differential Latent Attention (GDLA). Combines two ideas: 1. Differential Attention: cancels noise by subtracting a secondary attention map from the primary (see DiffAttentionV2). 2. Latent Attention (MLA-style, DeepSeek V3): compresses KV into a low-rank latent space to reduce KV-cache size during inference. - Q path: x -> Wq_a (down-proj to q_lora_rank) -> RMSNorm -> Wq_b (up-proj to n_heads * qk_head_dim) - KV path: x -> Wkv_a (down-proj to kv_lora_rank + qk_rope_head_dim) -> split -> RMSNorm -> Wkv_b Only the low-rank latent (kv_lora_rank dims) needs to be cached. Each QK head is split into nope (no positional encoding) and rope parts: qk_head_dim = qk_nope_head_dim + qk_rope_head_dim RoPE is applied only to the rope part. The rope component of K is shared across all KV heads (broadcast from a single head). A single KV latent with one norm; signal/noise separation happens after projection via grouped head splitting. """ def __init__(self, layer_idx, model_args: MotifModelArgs): super().__init__() self.dim = model_args.dim # SWA must be determined first so head counts can be overridden below. self.window_size = _compute_window_size(layer_idx, model_args) is_swa = self.window_size[0] > 0 # Head counts: SWA layers may use more heads than full-attention layers. swa_n_heads = getattr(model_args, "swa_n_heads", None) swa_n_kv_heads = getattr(model_args, "swa_n_kv_heads", None) if is_swa and swa_n_heads is not None: self.n_heads = swa_n_heads self.n_kv_heads = swa_n_kv_heads if swa_n_kv_heads is not None else swa_n_heads else: self.n_heads = model_args.n_heads self.n_kv_heads = model_args.n_heads if model_args.n_kv_heads is None else model_args.n_kv_heads self.q_lora_rank = model_args.q_lora_rank self.kv_lora_rank = model_args.kv_lora_rank self.head_dim = model_args.dim // model_args.n_heads if model_args.head_dim is None else model_args.head_dim self.qk_rope_head_dim = ( model_args.qk_rope_head_dim if model_args.qk_rope_head_dim is not None else self.head_dim // 2 ) self.qk_nope_head_dim = self.head_dim - self.qk_rope_head_dim self.qk_head_dim = self.head_dim self.v_head_dim = model_args.v_head_dim if model_args.v_head_dim is not None else self.head_dim self.grouped_ratio = model_args.grouped_ratio if self.n_heads % (self.grouped_ratio + 1) != 0: raise ValueError(f"n_heads={self.n_heads} must be divisible by grouped_ratio+1={self.grouped_ratio + 1}") self.noise_heads = self.n_heads // (self.grouped_ratio + 1) self.elementwise_attn_output_gate = model_args.elementwise_attn_output_gate self.n_signal_heads = self.n_heads - self.noise_heads # Query projection with LoRA self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False) self.q_norm = activation.layers.RMSNorm(self.q_lora_rank, eps=model_args.norm_eps) if self.elementwise_attn_output_gate: gate_extra_heads = self.n_signal_heads self.wq_b = nn.Linear(self.q_lora_rank, int(self.n_heads) * self.qk_head_dim, bias=False) self.wq_b_gate = nn.Linear(self.q_lora_rank, int(gate_extra_heads) * self.v_head_dim, bias=False) else: self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False) # Key-value projection with LoRA-style structure self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False) self.kv_norm = activation.layers.RMSNorm(self.kv_lora_rank, eps=model_args.norm_eps) self.wkv_b = nn.Linear( self.kv_lora_rank, self.n_kv_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, ) self.lambda_proj = nn.Linear( self.dim, self.grouped_ratio * self.noise_heads, bias=False, ) self.wo = nn.Linear( self.grouped_ratio * self.noise_heads * self.v_head_dim, self.dim, bias=False, ) self.softmax_scale = self.qk_head_dim**-0.5 # YaRN softmax-scale adjustment (DeepSeek-V3): only on full-attn layers, # since SWA layers don't use YaRN RoPE extension. max_seq_len = model_args.max_seq_len original_seq_len = model_args.original_seq_len if (not is_swa) and max_seq_len > original_seq_len: rope_factor = model_args.rope_factor mscale_val = model_args.mscale mscale_val = 0.1 * mscale_val * math.log(rope_factor) + 1.0 self.softmax_scale = self.softmax_scale * mscale_val * mscale_val if layer_idx == 0: logger.debug( f"grouped_ratio : {self.grouped_ratio} " f"head_dim : {self.head_dim} " f"qk_head_dim : {self.qk_head_dim} " f"v_head_dim : {self.v_head_dim} " f"q_heads : {self.n_heads} " f"q_lora_rank : {self.q_lora_rank} " f"kv_lora_rank : {self.kv_lora_rank}" ) _use_fa4 = getattr(model_args, "use_fa4", False) self.diff_attn = DiffAttentionV2( head_dim=self.qk_head_dim, n_heads=self.n_heads, n_kv_heads=self.n_kv_heads, grouped_ratio=self.grouped_ratio, use_fa3_directly=model_args.use_fa3_directly, use_sdpa=model_args.use_sdpa, v_head_dim=self.v_head_dim, use_fa4=_use_fa4, ) self.layer_idx = layer_idx def init_weights(self, init_std: float): linear_list = [self.wq_a, self.wq_b, self.wkv_a] if self.elementwise_attn_output_gate: linear_list.append(self.wq_b_gate) linear_list.append(self.wkv_b) nn.init.trunc_normal_(self.wo.weight, mean=0.0, std=init_std, a=-3 * init_std, b=3 * init_std) for linear in linear_list: nn.init.trunc_normal_(linear.weight, mean=0.0, std=0.02, a=-3 * 0.02, b=3 * 0.02) self.kv_norm.reset_parameters() self.q_norm.reset_parameters() nn.init.zeros_(self.lambda_proj.weight) def reorder_headdim_elements_rope(self, qk, bs, seqlen): qk = qk.view(bs, seqlen, -1, self.qk_rope_head_dim // 2, 2) qk = qk.transpose(3, 4) return qk.reshape(bs, seqlen, -1, self.qk_rope_head_dim) def forward( self, x: torch.Tensor, freqs_cis: torch.Tensor, attention_mask: torch.Tensor | None = None, max_seqlen: int | None = None, position_ids: torch.Tensor | None = None, qk_clip: bool = False, ): if isinstance(freqs_cis, DTensor): assert freqs_cis.placements == (Replicate(),) freqs_cis = freqs_cis.to_local() bsz, seqlen, _ = x.size() # Ulysses / Ring SP: x is per-rank seq-shard [B, S/tp, ...] but # freqs_cis is full [S, ...]. RoPE is position-dependent, so slice # freqs_cis to this rank's seq range before applying. For Ulysses the # post-a2a head-shard still carries the correct per-position angles # because RoPE was applied first. For Ring SP no a2a happens — Q/K/V # stay token-sharded with their own slice's positions. _ulysses_group = getattr(self, "_ulysses_group", None) _ring_attn_group = getattr(self, "_ring_attn_group", None) _seq_parallel_group = _ulysses_group if _ulysses_group is not None else _ring_attn_group if _seq_parallel_group is not None: _rank_in_tp = torch.distributed.get_rank(group=_seq_parallel_group) freqs_cis = freqs_cis[_rank_in_tp * seqlen : (_rank_in_tp + 1) * seqlen] # Query projection q_latent = self.q_norm(self.wq_a(x)) q = self.wq_b(q_latent) q = q.view(bsz, seqlen, -1, self.qk_head_dim) # Break the view-chain so fused_q_rope_inplace can modify q in place # without tripping autograd's "leaf-view in-place" check. Needed under # both standard TP (DTensor wq_b) and Ulysses (FSDP-managed plain # weight whose all-gathered output still carries autograd tracking). q = q.clone().contiguous() # In-place fused Q RoPE: modifies q[..., nope_dim:] from interleaved → contiguous # Eliminates the subsequent cat(q_nope, q_pe) copy (~415µs/layer saved). gate_score = None if self.elementwise_attn_output_gate: gate_score = self.wq_b_gate(q_latent) gate_score = gate_score.view(bsz, seqlen, -1, int(self.v_head_dim)) kv = self.wkv_a(x) kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) q_total = activation.fused_q_rope_inplace(q, freqs_cis, self.qk_nope_head_dim, self.qk_rope_head_dim) k_pe = apply_rotary_emb_single(k_pe.unsqueeze(2), freqs_cis) k_pe = self.reorder_headdim_elements_rope(k_pe, bsz, seqlen) kv = kv.contiguous() kv_latent = self.wkv_b(self.kv_norm(kv)) # -1 on n_kv_heads so the TP-sharded local size is inferred. kv_latent = kv_latent.view(bsz, seqlen, -1, self.qk_nope_head_dim + self.v_head_dim) k_full, v = activation.fused_kv_split_rope_cat( kv_latent, k_pe, self.qk_nope_head_dim, self.v_head_dim, self.qk_rope_head_dim ) lambda_full = self.lambda_proj(x) # Per-layer SP dispatch (when attn_sp_mode="ring" (default) both # groups are armed in parallelize.py): SWA layers (window_size[0] # > 0) → ring (cheap global early-exit once q-k distance exceeds # W_L). Dense layers (window_size[0] < 0) → Ulysses (no ring # early-exit to amortize, and one a2a + bwd a2a is cheaper than a # full ring rotation). attn_sp_mode="ulysses": only _ulysses_group # is set, so every layer goes through the a2a path. _is_swa_layer = self.window_size[0] > 0 _use_ring = _is_swa_layer and _ring_attn_group is not None _use_ulysses = (not _use_ring) and _ulysses_group is not None if _use_ulysses: _tp_size = _ulysses_group.size() q_total = _ulysses_a2a(q_total, _ulysses_group, scatter_dim=2, gather_dim=1) k_full = _ulysses_a2a(k_full, _ulysses_group, scatter_dim=2, gather_dim=1) v = _ulysses_a2a(v, _ulysses_group, scatter_dim=2, gather_dim=1) # lambda_full last dim = noise_heads * grouped_ratio; reshape to expose # noise_heads as a head dim, a2a, then collapse back. lambda_full = lambda_full.view(bsz, seqlen, -1, self.grouped_ratio) lambda_full = _ulysses_a2a(lambda_full, _ulysses_group, scatter_dim=2, gather_dim=1) lambda_full = lambda_full.reshape(bsz, seqlen * _tp_size, -1) if gate_score is not None: gate_score = _ulysses_a2a(gate_score, _ulysses_group, scatter_dim=2, gather_dim=1) output, v_stats = self.diff_attn( q_total, k_full, v, lambda_full=lambda_full, attention_mask=attention_mask, max_seqlen=max_seqlen, position_ids=position_ids, gate_score=gate_score, qk_clip=qk_clip, window_size=self.window_size, softmax_scale=self.softmax_scale, ring_attn_group=_ring_attn_group if _use_ring else None, ) # Ulysses: head-shard → token-shard via all-to-all. # output: [B, S, n_heads/tp, v_head_dim] → [B, S/tp, n_heads, v_head_dim]. # Ring SP: output is already token-shard, no a2a needed. if _use_ulysses: output = _ulysses_a2a(output, _ulysses_group, scatter_dim=1, gather_dim=2) output = output.view(bsz, seqlen, -1) return self.wo(output), v_stats class TransformerBlock(nn.Module): """ TransformerBlock Module Args: layer_id (int): Identifier for the layer. model_args (MotifModelArgs): Model configuration arguments. Attributes: n_heads (int): Number of attention heads. dim (int): Dimension size of the model. head_dim (int): Dimension size of each attention head. attention (Attention): Attention module. feed_forward (FeedForward): FeedForward module. layer_id (int): Identifier for the layer. attention_norm (RMSNorm): Layer normalization for attention output. ffn_norm (RMSNorm): Layer normalization for feedforward output. """ ATTENTION_CLS_MAP: dict[str, type] = { "gdla": GDLAttention, } def __init__(self, layer_id: int, model_args: MotifModelArgs | MotifMoEModelArgs): super().__init__() self.n_heads = model_args.n_heads self.dim = model_args.dim attn_cls = self.ATTENTION_CLS_MAP.get(model_args.attention_cls) if attn_cls is None: raise ValueError( f"Unknown attention_cls={model_args.attention_cls!r}, expected one of {list(self.ATTENTION_CLS_MAP)}" ) if layer_id == 0: logger.info(f"Using {attn_cls.__name__}") self.attention = attn_cls(layer_id, model_args) # True for SWA layers → select swa-theta RoPE table in Motif.forward. self.is_swa = _compute_window_size(layer_id, model_args)[0] > 0 # use MoE layer for every interleave_moe_layer_step FFN layers if not isinstance(model_args, MotifMoEModelArgs): self.moe_enabled = False else: # use MoE layer for every interleave_moe_layer_step FFN layers, # but the first n_dense_first_layers layers are always dense n_dense_first = model_args.n_dense_first_layers self.moe_enabled = ( layer_id >= n_dense_first and (layer_id + 1) % model_args.interleave_moe_layer_step == 0 if model_args.interleave_moe_layer_step != 0 else False ) if self.moe_enabled: moe_args = model_args.moe_args _alb_per_layer = getattr(moe_args, "aux_load_balance_loss_weight_per_layer", None) if _alb_per_layer and layer_id in _alb_per_layer: moe_args = replace(moe_args, aux_load_balance_loss_weight=float(_alb_per_layer[layer_id])) dim = model_args.dim hidden_dim = model_args.hidden_dim self.moe = MoE( moe_args, dim=dim, hidden_dim=hidden_dim, use_mxfp8=getattr(model_args, "use_mxfp8", False), layer_id=layer_id, ) else: _moe_args = getattr(model_args, "moe_args", None) _pn_sigmoid = getattr(_moe_args, "polynorm_sigmoid_weight", None) if _pn_sigmoid is None: _pn_sigmoid = getattr(model_args, "polynorm_sigmoid_weight", True) _pn_global = getattr(_moe_args, "polynorm_output_scale", None) if _pn_global is None: _pn_global = float(getattr(model_args, "polynorm_output_scale", 1.0)) _pn_per_layer = getattr(_moe_args, "polynorm_output_scale_per_layer", None) or getattr( model_args, "polynorm_output_scale_per_layer", None ) _pn_end = ( float(_pn_per_layer[layer_id]) if (_pn_per_layer and layer_id in _pn_per_layer) else float(_pn_global) ) # Compute initial scale consistent with MoE schedule state. _pn_start = float(getattr(_moe_args, "polynorm_output_scale_start", _pn_end) if _moe_args else _pn_end) _pn_warmup = int(getattr(_moe_args, "polynorm_output_scale_warmup_steps", 0) if _moe_args else 0) _pn_active = _pn_warmup > 0 and _pn_start != _pn_end _pn_init = _pn_start if _pn_active else _pn_end self.feed_forward = FeedForward( dim=model_args.dim, hidden_dim=model_args.intermediate_dim, polynorm_sigmoid_weight=_pn_sigmoid, polynorm_output_scale=_pn_init, ) # Store end scale so _propagate_polynorm_scale_to_dense can interpolate per-layer. self.feed_forward._polynorm_scale_end = _pn_end # MHC (Manifold-constrained Hyper-Connections) configuration self.mhc_enabled = getattr(model_args, "mhc_enabled", False) self.mhc_expansion_rate = getattr(model_args, "mhc_expansion_rate", 4) if self.mhc_enabled: _hpa_start = float(getattr(model_args, "mhc_h_post_alpha_start", 0.0)) _hpa_end = float(getattr(model_args, "mhc_h_post_alpha_end", 0.0)) _hpa_warmup = int(getattr(model_args, "mhc_h_post_alpha_warmup_steps", 0)) # Two MHC layers: one for attention, one for FFN self.mhc_attn = MHCLayer( expansion_rate=self.mhc_expansion_rate, num_dim=model_args.dim, identity_init=getattr(model_args, "mhc_identity_init", True), sinkhorn_iters=getattr(model_args, "mhc_sinkhorn_iters", 20), h_post_alpha_start=_hpa_start, h_post_alpha_end=_hpa_end, h_post_alpha_warmup_steps=_hpa_warmup, ) self.mhc_ffn = MHCLayer( expansion_rate=self.mhc_expansion_rate, num_dim=model_args.dim, identity_init=getattr(model_args, "mhc_identity_init", True), sinkhorn_iters=getattr(model_args, "mhc_sinkhorn_iters", 20), h_post_alpha_start=_hpa_start, h_post_alpha_end=_hpa_end, h_post_alpha_warmup_steps=_hpa_warmup, ) # Use simple RMSNorm (no fused residual) when using MHC self.ffn_norm = activation.layers.RMSNorm(model_args.dim, eps=model_args.norm_eps) else: self.mhc_attn = None self.mhc_ffn = None # FusedAddRMSNorm for standard residual path self.ffn_norm = activation.layers.FusedAddRMSNorm(model_args.dim, eps=model_args.norm_eps) self.layer_idx = layer_id self.attention_norm = activation.layers.RMSNorm(model_args.dim, eps=model_args.norm_eps) # Outlier-suppression penalty on ffn_out: w * ReLU(|y| - threshold)^2. # Two regimes: # - "Earlier" layers (layer_id < n_layers - 1): use ffn_clip_loss_weight / # ffn_clip_threshold, capped at the first 52 transformer layers. # - "Tail" layers (the last main layer + all MTP layers, i.e. # layer_id >= n_layers - 1): use ffn_clip_loss_weight_last / # ffn_clip_threshold_last with a separate setting. if layer_id < model_args.n_layers - 1: _ffn_clip_layer_ok = layer_id < 52 self.ffn_clip_loss_weight = float(model_args.ffn_clip_loss_weight) if _ffn_clip_layer_ok else 0.0 self.ffn_clip_threshold = float(model_args.ffn_clip_threshold) else: self.ffn_clip_loss_weight = float(model_args.ffn_clip_loss_weight_last) self.ffn_clip_threshold = float(model_args.ffn_clip_threshold_last) if model_args.depth_init: self.weight_init_std = 0.02 / (2 * (layer_id + 1)) ** 0.5 else: self.weight_init_std = 0.02 / (2 * model_args.n_layers) ** 0.5 def forward( self, x: torch.Tensor, freqs_cis: torch.Tensor, attention_mask: torch.Tensor | None = None, max_seqlen: int | None = None, position_ids: torch.Tensor | None = None, qk_clip: bool = False, accumulated_load_balance_loss: torch.Tensor | None = None, ): """ Perform a forward pass through the TransformerBlock. Args: x (torch.Tensor): Input tensor. - Standard: (batch, seq_len, dim) - MHC: (batch, seq_len, expansion_rate, dim) freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies. accumulated_load_balance_loss: Running sum of aux load balance losses from previous MoE layers, threaded through the layer stack. Returns: torch.Tensor: Output tensor after applying attention and feedforward layers. accumulated_load_balance_loss: Updated accumulated aux loss (or None). """ if self.mhc_enabled: return self._forward_with_mhc( x, freqs_cis, attention_mask, max_seqlen, position_ids, qk_clip, accumulated_load_balance_loss ) else: return self._forward_standard( x, freqs_cis, attention_mask, max_seqlen, position_ids, qk_clip, accumulated_load_balance_loss ) def _ffn_clip_loss(self, ffn_out: torch.Tensor) -> torch.Tensor: """Outlier suppression on ffn_out: ReLU(|y| - threshold)^2 averaged over (B, S, D). """ return _ffn_clip_loss_fn(ffn_out, float(self.ffn_clip_threshold), float(self.ffn_clip_loss_weight)) def _ffn_sublayer( self, n_out: torch.Tensor, accumulated_load_balance_loss: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """MoE/dense FFN plus aux-loss and ffn-clip accumulation (shared by both residual paths).""" if self.moe_enabled: ffn_out, aux_loss = self.moe(n_out) if aux_loss is not None: accumulated_load_balance_loss = ( aux_loss if accumulated_load_balance_loss is None else accumulated_load_balance_loss + aux_loss ) else: ffn_out = self.feed_forward(n_out) if self.ffn_clip_loss_weight > 0: fcl = self._ffn_clip_loss(ffn_out) accumulated_load_balance_loss = ( fcl if accumulated_load_balance_loss is None else accumulated_load_balance_loss + fcl ) return ffn_out, accumulated_load_balance_loss def _forward_standard( self, x: torch.Tensor, freqs_cis: torch.Tensor, attention_mask: torch.Tensor | None, max_seqlen: int | None, position_ids: torch.Tensor | None, qk_clip: bool, accumulated_load_balance_loss: torch.Tensor | None, ): """Standard residual path (original behavior).""" x_normed = self.attention_norm(x) attn_out, v = self.attention( x_normed, freqs_cis, attention_mask, max_seqlen, position_ids, qk_clip, ) n_out, h = self.ffn_norm(x, attn_out) if isinstance(h, DTensor): # At TP, feed_forward_output is a torch.Tensor with SP shard. # So we can just convert h to local without redistributing. h = h.to_local() ffn_out, accumulated_load_balance_loss = self._ffn_sublayer(n_out, accumulated_load_balance_loss) out = h + ffn_out return out, v, accumulated_load_balance_loss def _forward_with_mhc( self, x: torch.Tensor, freqs_cis: torch.Tensor, attention_mask: torch.Tensor | None, max_seqlen: int | None, position_ids: torch.Tensor | None, qk_clip: bool, accumulated_load_balance_loss: torch.Tensor | None, ): """ MHC residual path. x shape: (batch, seq_len, expansion_rate, dim) """ # Attention sublayer with MHC residual x_reduced, mhc_attn_ctx = self.mhc_attn.pre(x) x_reduced_normed = self.attention_norm(x_reduced) attn_out, v = self.attention( x_reduced_normed, freqs_cis, attention_mask, max_seqlen, position_ids, qk_clip, ) h = self.mhc_attn.post(x, attn_out, mhc_attn_ctx) # FFN sublayer with MHC residual h_reduced, mhc_ffn_ctx = self.mhc_ffn.pre(h) n_out = self.ffn_norm(h_reduced) if isinstance(n_out, DTensor): n_out = n_out.to_local() ffn_out, accumulated_load_balance_loss = self._ffn_sublayer(n_out, accumulated_load_balance_loss) out = self.mhc_ffn.post(h, ffn_out, mhc_ffn_ctx) return out, v, accumulated_load_balance_loss def init_weights(self, buffer_device: torch.device): for norm in (self.attention_norm, self.ffn_norm): norm.reset_parameters() self.attention.init_weights(self.weight_init_std) if self.mhc_enabled: self.mhc_attn.init_weights(self.mhc_attn.identity_init) self.mhc_ffn.init_weights(self.mhc_ffn.identity_init) if self.moe_enabled: self.moe.init_weights(self.weight_init_std, buffer_device) else: self.feed_forward.init_weights(self.weight_init_std) @dataclass class MotifOutput: pred: torch.Tensor """ Output of the lm_head layer. Shape: (B, S, vocab_size) B: batch size S: sequence length """ attn_logits_dict: dict[int, torch.Tensor] | None """ Intermediate attention statistics from each decoder block. - None if qk_clip is False. - Key: decoder block index (int) - Value: per-head maximum attention score. Shape: (NH,) NH: number of attention heads """ loss: torch.Tensor | None = None """ Pre-computed loss when MTP is enabled. None when MTP is disabled (loss is computed externally). """ load_balance_loss: torch.Tensor | None = None """ Accumulated complementary sequence-wise auxiliary load balance loss across all MoE layers. None when aux_load_balance_loss_weight == 0 or model has no MoE layers. Applied in the trainer with the gradient-only trick: loss + (lb_loss - lb_loss.detach()). """ mtp_depth_losses: list[torch.Tensor] | None = None """ Per-depth MTP losses (unweighted scalars), where index 0 = NTP base loss, index i = MTP head at depth i. Only set when MTP is enabled. Used for per-depth MTP loss logging in the trainer. """ def compute_mtp_loss( model: "Motif", hidden: torch.Tensor, labels: torch.Tensor, loss_fn, input_embeds: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, max_seqlen: int | None = None, unembed_and_loss=None, qk_clip: bool = False, attn_logits_dict: dict[int, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, list[torch.Tensor]]: """Compute MTP (Multi-Token Prediction) loss. Called from ``Motif.forward()``. Requires ``model._mtp_hidden`` to be set by the forward pass before this function is called. Output projection (``model.output``) is handled internally — callers always pass hidden states, never logits. Args: model: Motif model instance (must have ``_mtp_hidden`` set by forward). hidden: Hidden states from the last transformer layer. labels: Ground truth token ids. loss_fn: Loss function (e.g. cross-entropy), used when unembed_and_loss is None. attention_mask: Optional attention mask for MTP transformer layers. max_seqlen: Optional max sequence length for MTP transformer layers. unembed_and_loss: Optional fused unembed+loss function that handles both vocab projection and loss computation internally. qk_clip: Whether to collect QK attention logits in MTP layers. Returns: Tuple of (weighted average loss, per-depth losses). Per-depth losses: list of scalar tensors [depth0_loss, depth1_loss, ...] (detached). """ depth_losses: list[torch.Tensor] = [] h = model._mtp_hidden shifted_labels = labels # Base prediction loss (token_idx == 0) if unembed_and_loss is not None: is_multi_mtp = model.multi_token_preds > 0 final_loss = unembed_and_loss(hidden, labels, unshard=True, reshard=not is_multi_mtp) else: # model.output(hidden) -> (B, S, V) logits, inside FSDP forward so weights are gathered final_loss = loss_fn(model.output(hidden), labels) depth_losses.append(final_loss.detach()) # depth 0 = NTP base loss weight_sum = 1.0 next_embeds = input_embeds def _run_mtp_block(h_in: torch.Tensor, next_h: torch.Tensor, token_idx: int) -> torch.Tensor: """input_proj + transformer + norm for one MTP step.""" h_out = model.mtp_layer[str(token_idx)](torch.cat([h_in, next_h], dim=-1)) layer_idx = token_idx + model.n_layers - 1 layer = model.layers[str(layer_idx)] freqs = model.freqs_cis_swa if (layer.is_swa and model.freqs_cis_swa is not None) else model.freqs_cis h_out, v, _ = layer( h_out, freqs, attention_mask=attention_mask, max_seqlen=max_seqlen, position_ids=None, qk_clip=qk_clip ) if qk_clip and v is not None and attn_logits_dict is not None: attn_logits_dict[layer_idx] = v return model.mtp_hidden_norm[str(token_idx)](h_out) def _embed_next_h(step_embeds: torch.Tensor) -> torch.Tensor: """Compute embed-path next_h: norm(embed) only, no shift(h).""" normed = model.mtp_embed_norm[str(token_idx)](step_embeds) if isinstance(h, DTensor) and not isinstance(normed, DTensor): normed = DTensor.from_local(normed, h.device_mesh, h.placements, run_check=False) return normed for token_idx in range(1, model.multi_token_preds + 1): if next_embeds is not None: next_embeds = shift_left_and_pad(next_embeds, 1, 0) step_embeds = next_embeds shifted_labels = shift_left_and_pad(shifted_labels, 1, -100) is_last_mtp = token_idx == model.multi_token_preds h = _run_mtp_block(h, _embed_next_h(step_embeds), token_idx) if unembed_and_loss is not None: mtp_loss = unembed_and_loss(h, shifted_labels, unshard=False, reshard=is_last_mtp) else: mtp_loss = loss_fn(model.output(h), shifted_labels) depth_losses.append(mtp_loss.detach()) weight = model.mtp_loss_weight * (model.mtp_depth_decay**token_idx) final_loss = final_loss + weight * mtp_loss weight_sum += weight return final_loss / weight_sum, depth_losses class Motif(nn.Module, ModelProtocol): """ Transformer Module Args: model_args (MotifModelArgs): Model configuration arguments. Attributes: model_args (MotifModelArgs): Model configuration arguments. vocab_size (int): Vocabulary size. n_layers (int): Number of layers in the model. tok_embeddings (ParallelEmbedding): Token embeddings. layers (torch.nn.ModuleList): List of Transformer blocks. norm (RMSNorm): Layer normalization for the model output. output (ColumnParallelLinear): Linear layer for final output. freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies. """ def __init__(self, model_args: MotifModelArgs): super().__init__() self.model_args = model_args self.vocab_size = model_args.vocab_size self.n_layers = model_args.n_layers self.eos_id = model_args.eos_id # MHC configuration self.mhc_enabled = getattr(model_args, "mhc_enabled", False) self.mhc_expansion_rate = getattr(model_args, "mhc_expansion_rate", 4) logger.info(f"padding_idx : {self.eos_id}") self.tok_embeddings = nn.Embedding(model_args.vocab_size, model_args.dim, padding_idx=self.eos_id) # TODO persistent should be set to false, since this buffer can be recomputed. # However, we set it to true because of pytorch/pytorch#123411: torch.compile # does not correctly handle non-persistent buffers. self.register_buffer("freqs_cis", self._precompute_freqs_cis(), persistent=True) # Separate RoPE table for SWA layers (plain RoPE, no YaRN) when the # model mixes full / SWA layers. None when all layers are full-attn. swa_freqs_cis = self._precompute_freqs_cis_swa() if swa_freqs_cis is not None: self.register_buffer("freqs_cis_swa", swa_freqs_cis, persistent=True) else: self.freqs_cis_swa = None self.layers = torch.nn.ModuleDict() for layer_id in range(model_args.n_layers): self.layers[str(layer_id)] = TransformerBlock(layer_id, model_args) self.norm = activation.layers.RMSNorm(model_args.dim, eps=model_args.norm_eps) self.output = nn.Linear(model_args.dim, model_args.vocab_size, bias=False) # Aux load balance loss reduction across MoE layers moe_args = getattr(model_args, "moe_args", None) self.n_moe_layers = sum(1 for layer in self.layers.values() if getattr(layer, "moe_enabled", False)) self.aux_lb_loss_reduction = getattr(moe_args, "aux_load_balance_loss_reduction", "sum") # MTP (Multi-Token Prediction) configuration self.mtp_enabled = getattr(model_args, "mtp_enabled", False) self.multi_token_preds = getattr(model_args, "multi_token_preds", 1) self.mtp_loss_weight = getattr(model_args, "mtp_loss_weight", 1.0) self.mtp_depth_decay = getattr(model_args, "mtp_depth_decay", 0.95) if self.mtp_enabled: # MTP layers: each head has a dedicated TransformerBlock (add_mtp_layer) # plus an embed-norm. self.mtp_layer = torch.nn.ModuleDict() self.mtp_hidden_norm = torch.nn.ModuleDict() self.mtp_embed_norm = torch.nn.ModuleDict() # Force dense FFN, SWA, and no MHC for MTP blocks: # - interleave_moe_layer_step=0: always use dense FeedForward (no MoE) # - sliding_window_pattern="all": all MTP layers use sliding window attention # - mhc_enabled=False: MTP layers never use MHC if isinstance(model_args, MotifMoEModelArgs): dense_args = replace( model_args, interleave_moe_layer_step=0, sliding_window_pattern="all", mhc_enabled=False ) else: dense_args = replace(model_args, sliding_window_pattern="all", mhc_enabled=False) for mtp_layer_id in range(1, self.multi_token_preds + 1): self.layers[str(mtp_layer_id + self.n_layers - 1)] = TransformerBlock( mtp_layer_id + self.n_layers - 1, dense_args ) self.mtp_layer[str(mtp_layer_id)] = nn.Linear(2 * model_args.dim, model_args.dim, bias=False) self.mtp_hidden_norm[str(mtp_layer_id)] = activation.layers.RMSNorm( model_args.dim, eps=model_args.norm_eps ) self.mtp_embed_norm[str(mtp_layer_id)] = activation.layers.RMSNorm( model_args.dim, eps=model_args.norm_eps ) logger.info(f"MTP enabled with {self.multi_token_preds} prediction heads") else: self.mtp_layer = None self.mtp_hidden_norm = None self.mtp_embed_norm = None self.init_weights() def init_weights( self, buffer_device: torch.device | None = None, ): """ [Note: On ``init_weights`` vs. ``reset_parameters``] 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. """ buffer_device = buffer_device or self.freqs_cis.device with torch.device(buffer_device): self.freqs_cis = self._precompute_freqs_cis() if self.freqs_cis_swa is not None: self.freqs_cis_swa = self._precompute_freqs_cis_swa() if self.tok_embeddings is not None: nn.init.trunc_normal_(self.tok_embeddings.weight, mean=0.0, std=0.02, a=-0.06, b=0.06) for layer in self.layers.values(): if layer is not None: layer.init_weights(buffer_device=buffer_device) if self.norm is not None: self.norm.reset_parameters() final_out_std = self.model_args.dim**-0.5 cutoff_factor = 3 if self.output is not None: 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, ) # Initialize MTP layers if self.mtp_enabled: for mtp_layer_id in range(1, len(self.mtp_layer) + 1): nn.init.trunc_normal_( self.mtp_layer[str(mtp_layer_id)].weight, mean=0.0, std=final_out_std, a=-cutoff_factor * final_out_std, b=cutoff_factor * final_out_std, ) self.mtp_hidden_norm[str(mtp_layer_id)].reset_parameters() if self.mtp_embed_norm is not None: self.mtp_embed_norm[str(mtp_layer_id)].reset_parameters() def _rope_head_dim(self) -> int: # Use qk_rope_head_dim only for GDLA attention, otherwise fall back to head_dim if getattr(self.model_args, "attention_cls", None) == "gdla" and self.model_args.qk_rope_head_dim is not None: return self.model_args.qk_rope_head_dim return ( self.model_args.dim // self.model_args.n_heads if self.model_args.head_dim is None else self.model_args.head_dim ) def _precompute_freqs_cis(self) -> torch.Tensor: """RoPE table for full-attention layers. YaRN-extended when max_seq_len > original_seq_len, else plain RoPE.""" rope_head_dim = self._rope_head_dim() max_seq_len = self.model_args.max_seq_len original_seq_len = self.model_args.original_seq_len # Need to compute until at least the max token limit for generation # TODO: document why we removed the 2x # relaxing in our CP enablement PR if max_seq_len > original_seq_len: return precompute_freqs_cis_yarn( rope_head_dim, max_seq_len, self.model_args.rope_theta, original_seq_len, self.model_args.rope_factor, self.model_args.beta_fast, self.model_args.beta_slow, ) return precompute_freqs_cis(rope_head_dim, max_seq_len, self.model_args.rope_theta) def _precompute_freqs_cis_swa(self) -> torch.Tensor | None: """RoPE table for SWA layers: plain RoPE with swa_rope_theta, no YaRN. Returns None when the model has no SWA layers (then freqs_cis covers all layers).""" if getattr(self.model_args, "sliding_window", None) is None: return None rope_head_dim = self._rope_head_dim() return precompute_freqs_cis( rope_head_dim, self.model_args.max_seq_len, self.model_args.swa_rope_theta, ) def forward( self, tokens: torch.Tensor, max_seqlen: int | None = None, input_batch: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, qk_clip: bool = False, loss_fct: torch.nn.Module | None = None, labels: torch.Tensor | None = None, unembed_and_loss=None, ): """ Perform a forward pass through the Transformer model. When MTP is disabled, this performs standard next-token prediction and returns logits in ``MotifOutput.pred``. When MTP is enabled, the method computes losses for multiple future-token prediction heads internally and returns the averaged loss in ``MotifOutput.loss`` (with ``pred=None``). MTP (``add_mtp_layer`` / ``embed`` mode): for each prediction step k > 0, the shifted input embeddings are normed (``mtp_embed_norm[k]``) and concatenated with the running hidden states, projected through a per-head linear layer (``mtp_layer[k]``), then run through a dedicated ``TransformerBlock`` (stored in ``self.layers`` at indices ``n_layers``, ``n_layers+1``, …) and normed by ``mtp_hidden_norm[k]``. The loss for head k is weighted by ``mtp_loss_weight * mtp_depth_decay ** k`` (base head weight 1.0); the final loss is their weight_sum-normalized weighted average. Args: tokens (torch.Tensor): Input token indices, shape (batch, seq_len). max_seqlen (int | None): Maximum sequence length for attention computation. input_batch (torch.Tensor): The input batch read from the dataloader (currently unused; accepted for interface compatibility). attention_mask (torch.Tensor | None): Attention mask tensor. qk_clip (bool): Whether to collect QK attention logits for clipping. loss_fct (torch.nn.Module | None): Loss function, required when MTP is enabled. labels (torch.Tensor | None): Ground truth labels, required when MTP is enabled. Returns: MotifOutput: When MTP is disabled, ``pred`` contains logits and ``loss`` is None. When MTP is enabled, ``pred`` is None and ``loss`` contains the averaged multi-token prediction loss. """ if qk_clip: attn_logits_dict = {} else: attn_logits_dict = None input_embeds = self.tok_embeddings(tokens) raw_embeds = input_embeds # pre-MHC, for MTP embed path # Expand to (B, S, E, D) for MHC if self.mhc_enabled: # Expand embedding output by repeating along new expansion dimension # (B, S, D) -> (B, S, E, D) input_embeds = input_embeds.unsqueeze(2).expand(-1, -1, self.mhc_expansion_rate, -1).contiguous() h = input_embeds accumulated_load_balance_loss = None # Only run the first n_layers base layers. For add_mtp_layer, extra MTP # TransformerBlocks live in self.layers at indices >= n_layers and must # NOT be run here; they are invoked per-head in the MTP loop below. base_layers = [v for k, v in self.layers.items() if int(k) < self.n_layers] for idx, layer in enumerate(base_layers): layer_freqs_cis = ( self.freqs_cis_swa if (layer.is_swa and self.freqs_cis_swa is not None) else self.freqs_cis ) h, v, accumulated_load_balance_loss = layer( h, layer_freqs_cis, attention_mask=attention_mask, max_seqlen=max_seqlen, position_ids=None, qk_clip=qk_clip, accumulated_load_balance_loss=accumulated_load_balance_loss, ) if qk_clip: assert v is not None attn_logits_dict[idx] = v # Reduce from (B, S, E, D) back to (B, S, D) for MHC if self.mhc_enabled: # Take mean across expansion dimension h = h.mean(dim=2) if self.norm is not None: h = self.norm(h) if accumulated_load_balance_loss is not None and self.aux_lb_loss_reduction == "mean" and self.n_moe_layers > 1: accumulated_load_balance_loss = accumulated_load_balance_loss / self.n_moe_layers # Plain next-token-prediction path when MTP is disabled. if self.mtp_hidden_norm is None: if unembed_and_loss is not None and labels is not None: loss = unembed_and_loss(h, labels) return MotifOutput( pred=None, attn_logits_dict=attn_logits_dict, loss=loss, load_balance_loss=accumulated_load_balance_loss, ) output = self.output(h) return MotifOutput( pred=output, attn_logits_dict=attn_logits_dict, loss=None, load_balance_loss=accumulated_load_balance_loss, ) assert loss_fct is not None, "MTP enabled: loss_fct must be provided" assert labels is not None, "MTP enabled: labels must be provided" # Compute MTP loss (base prediction + MTP heads) self._mtp_hidden = h final_loss, mtp_depth_losses = compute_mtp_loss( model=self, hidden=h, labels=labels, loss_fn=loss_fct, input_embeds=raw_embeds, attention_mask=attention_mask, max_seqlen=max_seqlen, unembed_and_loss=unembed_and_loss, qk_clip=qk_clip, attn_logits_dict=attn_logits_dict, ) return MotifOutput( pred=None, attn_logits_dict=attn_logits_dict, loss=final_loss, load_balance_loss=accumulated_load_balance_loss, mtp_depth_losses=mtp_depth_losses, )