From 5963693aaaa305401ad7f3ea961b7bef0875d957 Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis Date: Sun, 23 Aug 2026 20:12:22 -0700 Subject: [PATCH 1/3] feat(laguna): support packed THD context parallelism Signed-off-by: Alexandros Koumparoulis --- docs/model-coverage/llm/poolside/laguna.mdx | 9 +- .../laguna/laguna_xs_2p1_hellaswag_ep8.yaml | 4 +- .../laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml | 139 +++++++++++ .../distributed/blockdiag_cp/batch.py | 46 +++- .../distributed/blockdiag_cp/kernels.py | 11 +- .../distributed/blockdiag_cp/runtime.py | 39 ++- .../components/models/laguna/model.py | 228 +++++++++++++++--- .../distributed/test_blockdiag_cp.py | 53 +++- tests/unit_tests/models/laguna/test_model.py | 195 +++++++++++++++ 9 files changed, 673 insertions(+), 51 deletions(-) create mode 100644 examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml diff --git a/docs/model-coverage/llm/poolside/laguna.mdx b/docs/model-coverage/llm/poolside/laguna.mdx index 8de7d51b5c..c5fd539129 100644 --- a/docs/model-coverage/llm/poolside/laguna.mdx +++ b/docs/model-coverage/llm/poolside/laguna.mdx @@ -3,7 +3,7 @@ title: "Laguna" description: "" slug: model-coverage/large-language-models/poolside/laguna --- -[Laguna](https://huggingface.co/poolside) is Poolside's hybrid-attention MoE language model family. The Automodel implementation supports full SFT with expert parallelism for Laguna S 2.1 and Laguna XS 2.1. +[Laguna](https://huggingface.co/poolside) is Poolside's hybrid-attention MoE language model family. The Automodel implementation supports full SFT with expert parallelism for Laguna S 2.1 and Laguna XS 2.1, including native THD sequence packing with context parallelism. @@ -21,6 +21,7 @@ slug: model-coverage/large-language-models/poolside/laguna - `LagunaForCausalLM` - Layer-specific attention head counts, QK RMSNorm, and softplus attention output gating. - Full and sliding-window attention layers can use separate RoPE settings. +- Packed THD attention preserves document and sliding-window boundaries with block-diagonal context parallelism. - MoE blocks use `nemo_automodel.components.moe.layers.MoE` with sigmoid routing, top-k probability normalization, fp32 gate compute, correction bias loading, one shared expert, and grouped expert weights for EP. ## Example Recipes @@ -29,6 +30,7 @@ slug: model-coverage/large-language-models/poolside/laguna |---|---| | [laguna_s_2p1_hellaswag_ep16.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_s_2p1_hellaswag_ep16.yaml) | SFT — Laguna S 2.1 on HellaSwag with EP16 | | [laguna_xs_2p1_hellaswag_ep8.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml) | SFT — Laguna XS 2.1 on HellaSwag with EP8 and 1K NEAT sequence packing | +| [laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml) | SFT — Laguna XS 2.1 on HellaSwag with EP8, CP2, and 1K THD sequence packing | ## Run the Recipe @@ -36,8 +38,11 @@ slug: model-coverage/large-language-models/poolside/laguna # Laguna XS 2.1 on one 8-GPU node uv run automodel --nproc-per-node=8 examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml +# Laguna XS 2.1 with THD packing and CP2 on one 8-GPU node +uv run automodel --nproc-per-node=8 examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml + # Laguna S 2.1 on two 8-GPU nodes uv run automodel --nproc-per-node=8 examples/llm_finetune/laguna/laguna_s_2p1_hellaswag_ep16.yaml ``` -The Laguna XS 2.1 recipe uses `ep_size: 8` on one 8-GPU node and packs documents into 1,024-token sequences with `packing_strategy: neat`. Laguna currently supports expert parallelism, but not native THD packing or context parallelism; keep `cp_size: 1`. For the Laguna S 2.1 recipe, submit through your cluster launcher with two 8-GPU nodes for the default `ep_size: 16`. +The baseline Laguna XS 2.1 recipe uses `ep_size: 8` and NEAT packing with `cp_size: 1`. The THD variant uses block-diagonal SDPA, `packing_strategy: thd`, and `cp_size: 2` while retaining EP8 on the same eight GPUs. For Laguna S 2.1, submit through your cluster launcher with two 8-GPU nodes for the default `ep_size: 16`. diff --git a/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml b/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml index 7e543f686a..34f45336ad 100644 --- a/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml +++ b/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml @@ -18,8 +18,8 @@ # automodel examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml --nproc-per-node 8 # # EP size must divide num_experts (256). ep_size=8 -> 32 experts/rank. -# Laguna currently uses SDPA with NEAT packing. Native THD packing and context -# parallelism are not supported by this model implementation. +# This baseline uses SDPA with NEAT packing and no context parallelism. See +# laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml for native THD packing with CP2. recipe: TrainFinetuneRecipeForNextTokenPrediction diff --git a/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml b/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml new file mode 100644 index 0000000000..e48d9009e7 --- /dev/null +++ b/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml @@ -0,0 +1,139 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +# Laguna XS 2.1 SFT with native THD packing, CP2, EP8, and FSDP2. +# +# Hardware target: 1 node x 8 GPUs. +# automodel examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml --nproc-per-node 8 +# +# CP2 shards every packed document across two ranks with block-diagonal SDPA. +# EP8 partitions the checkpoint's 256 routed experts into 32 experts per rank. + +recipe: TrainFinetuneRecipeForNextTokenPrediction + +step_scheduler: + global_batch_size: 8 + local_batch_size: 1 + ckpt_every_steps: 200 + val_every_steps: 50 + num_epochs: 1 + max_steps: 100 + +dist_env: + backend: nccl + timeout_minutes: 60 + +rng: + _target_: nemo_automodel.components.training.rng.StatefulRNG + seed: 1111 + ranked: true + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: poolside/Laguna-XS-2.1 + trust_remote_code: false + torch_dtype: bfloat16 + attn_implementation: sdpa + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch_fp32 + experts: torch_mm + dispatcher: torch + fake_balanced_gate: false + gate_precision: float32 + enable_hf_state_dict_adapter: true + enable_fsdp_optimizations: true + +checkpoint: + enabled: false + checkpoint_dir: checkpoints/laguna_xs_2p1_hellaswag_ep8_cp2_thd + model_save_format: safetensors + save_consolidated: final + +distributed: + strategy: fsdp2 + tp_size: 1 + cp_size: 2 + pp_size: 1 + ep_size: 8 + + sequence_parallel: false + activation_checkpointing: true + + moe: + reshard_after_forward: false + wrap_outer_model: false + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +dataset: + _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag + path_or_dataset: rowan/hellaswag + split: train + pad_to_max_length: false + tokenizer: + _target_: transformers.AutoTokenizer.from_pretrained + pretrained_model_name_or_path: poolside/Laguna-XS-2.1 + trust_remote_code: false + +packed_sequence: + packed_sequence_size: 1024 + packing_strategy: thd + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: nemo_automodel.components.datasets.utils.packed_sequence_thd_collater + shuffle: true + num_workers: 4 + +validation_dataset: + _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag + path_or_dataset: rowan/hellaswag + split: validation + num_samples_limit: 128 + pad_to_max_length: false + tokenizer: + _target_: transformers.AutoTokenizer.from_pretrained + pretrained_model_name_or_path: poolside/Laguna-XS-2.1 + trust_remote_code: false + +validation_dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: nemo_automodel.components.datasets.utils.packed_sequence_thd_collater + shuffle: false + drop_last: true + num_workers: 4 + +optimizer: + _target_: torch.optim.AdamW + betas: [0.9, 0.95] + eps: 1e-8 + lr: 1e-5 + weight_decay: 0.0 + +wandb: + enable: false + project: laguna-xs-2p1-sft + name: laguna_xs_2p1_hellaswag_ep8_cp2_thd_1k + mode: online + +ci: + recipe_owner: akoumpa + # cp_size(2) and ep_size(8) share one 8-GPU node. + nodes: 1 + nproc_per_node: 8 + time: "01:00:00" diff --git a/nemo_automodel/components/distributed/blockdiag_cp/batch.py b/nemo_automodel/components/distributed/blockdiag_cp/batch.py index eb06316b7d..c725c7fffb 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/batch.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/batch.py @@ -31,9 +31,11 @@ def _cp_blockdiag_doc_ids(batch: dict, seq_len: int, device, batch_size: int) -> """Resolve per-position document ids ``[B, S]`` (0 == padding) for the mask. Prefers the collator's ``_packed_seq_ids`` (1-based document index per token, - present when a pack holds >1 document). Otherwise falls back to the 4-D - block-causal ``attention_mask`` diagonal (valid positions) or, lacking both, - treats the whole sequence as a single document. + present when a pack holds >1 document). THD batches instead carry + ``seq_lens`` and ``seq_lens_padded``; those are expanded into document ids + while keeping inter-document padding at id 0. Otherwise falls back to the + 4-D block-causal ``attention_mask`` diagonal (valid positions) or, lacking + both, treats the whole sequence as a single document. Args: batch: The training batch; may contain ``_packed_seq_ids`` ``[B, S]`` @@ -49,6 +51,44 @@ def _cp_blockdiag_doc_ids(batch: dict, seq_len: int, device, batch_size: int) -> seq_ids = batch.get("_packed_seq_ids", None) if seq_ids is not None: return seq_ids.to(device=device, dtype=torch.long) + seq_lens = batch.get("seq_lens") + if isinstance(seq_lens, torch.Tensor): + padded_lens = batch.get("seq_lens_padded", seq_lens) + if not isinstance(padded_lens, torch.Tensor): + raise ValueError("THD block-diagonal CP requires tensor seq_lens_padded metadata.") + if seq_lens.ndim == 1: + seq_lens = seq_lens.unsqueeze(0) + if padded_lens.ndim == 1: + padded_lens = padded_lens.unsqueeze(0) + if seq_lens.shape[0] != batch_size or padded_lens.shape[0] != batch_size: + raise ValueError( + "THD block-diagonal CP sequence metadata batch dimension must match input_ids: " + f"seq_lens={tuple(seq_lens.shape)}, seq_lens_padded={tuple(padded_lens.shape)}, " + f"batch_size={batch_size}." + ) + + rows = [] + for row_idx in range(batch_size): + pieces = [] + document_id = 1 + for actual, padded in zip(seq_lens[row_idx].tolist(), padded_lens[row_idx].tolist()): + if actual < 0 or padded < 0: + continue + if actual > padded: + raise ValueError(f"THD sequence length {actual} exceeds padded length {padded}.") + pieces.append(torch.full((actual,), document_id, dtype=torch.long, device=device)) + if padded > actual: + pieces.append(torch.zeros(padded - actual, dtype=torch.long, device=device)) + document_id += 1 + row = torch.cat(pieces) if pieces else torch.empty(0, dtype=torch.long, device=device) + if row.numel() > seq_len: + raise ValueError( + f"THD sequence metadata covers {row.numel()} tokens, exceeding sequence length {seq_len}." + ) + if row.numel() < seq_len: + row = torch.cat([row, torch.zeros(seq_len - row.numel(), dtype=torch.long, device=device)]) + rows.append(row) + return torch.stack(rows) attn = batch.get("attention_mask", None) if attn is not None and attn.dim() == 4: # [B, 1, S, S] block-causal bool -> diagonal gives per-position validity. diff --git a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py index b95901ddd3..920cb0eee1 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/kernels.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/kernels.py @@ -361,8 +361,9 @@ def _cp_blockdiag_mask( local_len: int, full_len: int, batch_size: int, + window_size: tuple[int, int] | None = None, ) -> torch.Tensor: - """Per-document causal attention mask for block-diagonal CP, shape ``[B, 1, L, S]``. + """Per-document causal or sliding mask for block-diagonal CP, shape ``[B, 1, L, S]``. ``doc_ids`` is the full (all-rank, padded) per-position document index ``[B, S]`` (0 == padding). Query rows are this rank's local positions @@ -381,6 +382,8 @@ def _cp_blockdiag_mask( local_len: ``L``, the number of local query rows. full_len: ``S``, the number of key columns (full padded sequence). batch_size: ``B``, used to expand a 1D ``doc_ids``. + window_size: Optional inclusive ``(left, right)`` local-attention window. + Negative values leave that side unbounded. Returns: Boolean allow-mask ``[B, 1, L, S]`` (True == may attend). @@ -396,6 +399,12 @@ def _cp_blockdiag_mask( row_pos = torch.arange(row_offset, row_offset + L, device=device).view(1, L, 1) col_pos = torch.arange(S, device=device).view(1, 1, S) causal = row_pos >= col_pos # [1, L, S] + if window_size is not None: + left, right = window_size + if left is not None and left >= 0: + causal = causal & (col_pos >= row_pos - left) + if right is not None and right > 0: + causal = causal & (col_pos <= row_pos + right) # Always allow the diagonal (q_pos == k_pos) so every query attends to >=1 key even # in all-pad/empty rows -- prevents NaN/hang. self_diag = row_pos == col_pos # [1, L, S] diff --git a/nemo_automodel/components/distributed/blockdiag_cp/runtime.py b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py index 3fac33fa4f..a3d3137003 100644 --- a/nemo_automodel/components/distributed/blockdiag_cp/runtime.py +++ b/nemo_automodel/components/distributed/blockdiag_cp/runtime.py @@ -223,6 +223,7 @@ def cp_blockdiag_sdpa( is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, + window_size: tuple[int, int] | None = None, **kwargs, ) -> torch.Tensor: """Block-diagonal context-parallel SDPA. @@ -247,6 +248,8 @@ def cp_blockdiag_sdpa( is_causal: Ignored on the CP path (forwarded to stock SDPA otherwise). scale: Softmax scale (``None`` -> ``D**-0.5``). enable_gqa: Grouped-query attention flag as passed by HF's sdpa path. + window_size: Optional inclusive ``(left, right)`` sliding-attention + window. Negative values leave that side unbounded. **kwargs: Ignored; accepted for SDPA signature compatibility. Returns: @@ -280,16 +283,21 @@ def cp_blockdiag_sdpa( # Decide the KV-exchange path (needed-only halo/a2a vs full all-gather) explicitly, # logging why, including downgrades to all-gather for mode, kernel, missing meta, or # a cross-node CP group where needed-only exchange is disabled by default. - path, plan, reason = _select_kv_exchange_path( - step_state, - group, - doc_ids, - query.shape[seq_dim], - query.device, - offset, - query_dtype=query.dtype, - dropout_p=dropout_p, - ) + left_window, right_window = window_size or (-1, 0) + has_local_window = (left_window is not None and left_window >= 0) or (right_window is not None and right_window > 0) + if has_local_window: + path, plan, reason = "allgather", None, "sliding-window mask requires the dense block-diagonal path" + else: + path, plan, reason = _select_kv_exchange_path( + step_state, + group, + doc_ids, + query.shape[seq_dim], + query.device, + offset, + query_dtype=query.dtype, + dropout_p=dropout_p, + ) global _KV_EXCHANGE_PATH_LOGGED if not _KV_EXCHANGE_PATH_LOGGED: _KV_EXCHANGE_PATH_LOGGED = True @@ -357,7 +365,7 @@ def cp_blockdiag_sdpa( key_full = kv_full[:, :n_kv_heads_local] value_full = kv_full[:, n_kv_heads_local:] - if attn_backend in ("flash", "te"): + if not has_local_window and attn_backend in ("flash", "te"): out = kernels._cp_blockdiag_varlen( query, key_full, @@ -392,7 +400,14 @@ def cp_blockdiag_sdpa( L = query.shape[seq_dim] S = key_full.shape[seq_dim] - allow = kernels._cp_blockdiag_mask(doc_ids, offset, L, S, B) # [B, 1, L, S] + allow = kernels._cp_blockdiag_mask( + doc_ids, + offset, + L, + S, + B, + window_size=window_size if has_local_window else None, + ) # [B, 1, L, S] return _ORIGINAL_SDPA( query, diff --git a/nemo_automodel/components/models/laguna/model.py b/nemo_automodel/components/models/laguna/model.py index c677927fa2..73f5ee44b2 100644 --- a/nemo_automodel/components/models/laguna/model.py +++ b/nemo_automodel/components/models/laguna/model.py @@ -17,6 +17,7 @@ import copy from collections.abc import Callable from dataclasses import dataclass +from functools import partial from typing import Any, Union import torch @@ -26,6 +27,15 @@ from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from nemo_automodel.components.attention.utils import ( + initialize_attn_module_and_func, + postprocess_output_for_attn, + preprocess_args_and_kwargs_for_attn, +) +from nemo_automodel.components.distributed.context_parallel.sharder import ( + ContextParallelSharder, + contiguous_local_indices, +) from nemo_automodel.components.models.common import BackendConfig, initialize_linear_module from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin from nemo_automodel.components.models.common.tie_word_embeddings import ( @@ -42,6 +52,7 @@ from nemo_automodel.components.moe.config import MoEConfig from nemo_automodel.components.moe.fsdp_mixin import MoEFSDPSyncMixin from nemo_automodel.components.moe.layers import MLP, MoE +from nemo_automodel.components.utils.model_utils import squeeze_input_for_thd from nemo_automodel.shared.utils import dtype_from_str as get_dtype @@ -68,10 +79,13 @@ def _apply_rotary_pos_emb( """Apply RoPE to query and key states. Args: - q: Query tensor of shape [batch, heads, sequence, head_dim]. - k: Key tensor of shape [batch, key_value_heads, sequence, head_dim]. - cos: Cosine tensor of shape [batch, sequence, rotary_dim]. - sin: Sine tensor of shape [batch, sequence, rotary_dim]. + q: Query tensor of shape [batch, heads, sequence, head_dim] for BSHD or + [total_tokens, heads, head_dim] for THD. + k: Key tensor of shape [batch, key_value_heads, sequence, head_dim] for + BSHD or [total_tokens, key_value_heads, head_dim] for THD. + cos: Cosine tensor of shape [batch, sequence, rotary_dim] for BSHD or + [total_tokens, rotary_dim] for THD. + sin: Sine tensor with the same shape as ``cos``. Returns: Tuple of rotated query and key tensors with the same shapes as ``q`` and ``k``. @@ -364,12 +378,18 @@ def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Te """Compute RoPE cosine and sine tables. Args: - x: Activation tensor of shape [batch, sequence, hidden], used for device and dtype. - position_ids: Position tensor of shape [batch, sequence]. + x: Activation tensor of shape [batch, sequence, hidden] for BSHD or + [total_tokens, hidden] for THD, used for device and dtype. + position_ids: Position tensor of shape [batch, sequence] for BSHD or + [total_tokens] for THD. Returns: - Tuple of cosine and sine tensors, each of shape [batch, sequence, rotary_dim]. + Tuple of cosine and sine tensors shaped [batch, sequence, rotary_dim] + for BSHD or [total_tokens, rotary_dim] for THD. """ + is_thd = position_ids.ndim == 1 + if is_thd: + position_ids = position_ids.unsqueeze(0) inv_freq = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids = position_ids[:, None, :].float() device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" @@ -378,7 +398,12 @@ def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Te emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + cos = cos.to(dtype=x.dtype) + sin = sin.to(dtype=x.dtype) + if is_thd: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + return cos, sin class LagunaAttention(nn.Module): @@ -449,6 +474,19 @@ def __init__(self, config: LagunaConfig, backend: BackendConfig, layer_idx: int) self.q_norm = LagunaRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) self.k_norm = LagunaRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) + if backend.attn == "te": + # Ordinary BSHD inputs retain Laguna's HuggingFace-compatible attention + # path. TE is selected only for native packed THD batches. + self._te_thd_only = True + self.attn_module, self.attn_func = initialize_attn_module_and_func( + attn_impl="te", + num_attention_heads=self.num_heads, + num_qk_channels=self.head_dim, + num_v_channels=self.head_dim, + softmax_scale=self.scaling, + num_gqa_groups=self.num_key_value_heads, + attention_dropout=self.attention_dropout, + ) def forward( self, @@ -460,17 +498,80 @@ def forward( """Run Laguna attention for one decoder layer. Args: - hidden_states: Tensor of shape [batch, sequence, hidden]. + hidden_states: Tensor of shape [batch, sequence, hidden] for BSHD or + [total_tokens, hidden] for packed THD. position_embeddings: Tuple of cosine and sine RoPE tensors, each of shape - [batch, sequence, rotary_dim]. + [batch, sequence, rotary_dim] for BSHD or [total_tokens, rotary_dim] + for THD. attention_mask: Optional additive mask broadcastable to [batch, heads, sequence, key_sequence]. - **kwargs: Additional attention backend arguments. + **kwargs: Additional attention backend arguments. THD requires + ``qkv_format='thd'`` and ``cu_seqlens``; CP additionally supplies + ``cp_size`` and ``cp_rank``. Returns: - Tuple of attention output [batch, sequence, hidden] and optional attention weights - [batch, heads, sequence, key_sequence]. + Tuple of attention output shaped like ``hidden_states`` and optional BSHD + attention weights. THD returns ``None`` for the weights. """ + if kwargs.get("qkv_format") == "thd": + if hidden_states.ndim != 2: + raise ValueError(f"THD attention requires hidden_states [T, H], got {tuple(hidden_states.shape)}.") + + token_count = hidden_states.shape[0] + query_states = self.q_proj(hidden_states).view(token_count, self.num_heads, self.head_dim) + key_states = self.k_proj(hidden_states).view(token_count, self.num_key_value_heads, self.head_dim) + value_states = self.v_proj(hidden_states).view(token_count, self.num_key_value_heads, self.head_dim) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + cos, sin = position_embeddings + query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin) + + window_size = (-1, 0) if self.sliding_window is None else (self.sliding_window - 1, 0) + from nemo_automodel.components.distributed.blockdiag_cp import ( + cp_blockdiag_sdpa, + current_blockdiag_cp_state, + ) + + if current_blockdiag_cp_state() is not None: + if self.backend.attn != "sdpa": + raise ValueError("Laguna packed context parallelism requires backend.attn='sdpa'.") + attn_output = cp_blockdiag_sdpa( + query_states.transpose(0, 1).unsqueeze(0), + key_states.transpose(0, 1).unsqueeze(0), + value_states.transpose(0, 1).unsqueeze(0), + dropout_p=0.0 if not self.training else self.attention_dropout, + scale=self.scaling, + enable_gqa=True, + window_size=window_size, + ) + attn_output = attn_output.squeeze(0).transpose(0, 1).contiguous().flatten(1) + else: + if getattr(self, "attn_module", None) is None: + raise ValueError( + "Packed THD attention requires backend.attn='te', or backend.attn='sdpa' with context parallelism." + ) + query_states, key_states, value_states, te_kwargs = preprocess_args_and_kwargs_for_attn( + query_states, + key_states, + value_states, + attention_mask, + "te", + window_size=window_size, + **kwargs, + ) + attn_output = self.attn_func(query_states, key_states, value_states, **te_kwargs) + attn_output = postprocess_output_for_attn(attn_output, "te").flatten(1) + + if self.g_proj is not None: + gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype) + if self.gating_mode == "per-head": + attn_output = ( + attn_output.view(token_count, self.num_heads, self.head_dim) * gate.unsqueeze(-1) + ).flatten(1) + else: + attn_output = attn_output * gate + return self.o_proj(attn_output), None + batch, seq_len = hidden_states.shape[:2] query_states = self.q_proj(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = self.k_proj(hidden_states).view( @@ -526,6 +627,12 @@ def forward( return self.o_proj(attn_output), attn_weights + def setup_cp_attention(self, cp_mesh) -> None: + """Record that Laguna uses its model-owned packed block-diagonal CP path.""" + del cp_mesh + if self.backend.attn != "sdpa": + raise ValueError("Laguna packed context parallelism requires backend.attn='sdpa'.") + def init_weights(self, buffer_device: torch.device, init_std: float = 0.02) -> None: del buffer_device for linear in (self.q_proj, self.k_proj, self.v_proj, self.o_proj, self.g_proj): @@ -784,33 +891,48 @@ def forward( """Run the Laguna decoder stack. Args: - input_ids: Optional token IDs of shape [batch, sequence]. - inputs_embeds: Optional embeddings of shape [batch, sequence, hidden]. - position_ids: Optional position IDs of shape [batch, sequence]. + input_ids: Optional token IDs of shape [batch, sequence] for BSHD or + [total_tokens] for packed THD. + inputs_embeds: Optional embeddings of shape [batch, sequence, hidden] + for BSHD or [total_tokens, hidden] for THD. + position_ids: Optional position IDs of shape [batch, sequence] for BSHD + or [total_tokens] for THD. attention_mask: Optional 2D bool/int mask of shape [batch, sequence], a 4D additive mask, or a mapping with per-attention-type masks keyed by "full_attention" and "sliding_attention". padding_mask: Optional bool tensor of shape [batch, sequence], where True marks tokens excluded from MoE routing. - **kwargs: Additional attention backend arguments. + **kwargs: Additional attention backend arguments. Packed THD sets + ``qkv_format='thd'`` and carries cumulative document lengths. Returns: - Final hidden states of shape [batch, sequence, hidden]. + Final hidden states shaped like ``inputs_embeds``. """ if inputs_embeds is None: if input_ids is None: raise ValueError("input_ids or inputs_embeds must be provided") inputs_embeds = self.embed_tokens(input_ids) + is_thd = kwargs.get("qkv_format") == "thd" + if is_thd and inputs_embeds.ndim != 2: + raise ValueError(f"THD model input must be [T, H], got {tuple(inputs_embeds.shape)}.") if position_ids is None: - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + if is_thd: + if kwargs.get("cp_size", 1) > 1: + raise ValueError("THD context parallelism requires explicit position_ids.") + position_ids = torch.arange(inputs_embeds.shape[0], device=inputs_embeds.device) + else: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) if padding_mask is None and isinstance(attention_mask, torch.Tensor): padding_mask = _derive_padding_mask(attention_mask) - causal_mask_mapping = self._build_causal_mask_mapping( - inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - ) + if is_thd: + causal_mask_mapping = {"full_attention": None, "sliding_attention": None} + else: + causal_mask_mapping = self._build_causal_mask_mapping( + inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) hidden_states = inputs_embeds full_position_embeddings = self.rotary_emb(hidden_states, position_ids) @@ -854,6 +976,8 @@ class LagunaForCausalLM(HFCheckpointingMixin, nn.Module, MoEFSDPSyncMixin): """Causal LM wrapper for Laguna with Automodel checkpoint adapters.""" tie_word_embeddings_support: TieSupport = TieSupport.UNTIED_ONLY + _supports_cp_sdpa = True + _packed_cp_attn_backends = ("sdpa",) _keep_in_fp32_modules_strict = ["mlp.gate.e_score_correction_bias", "rotary_emb"] _skip_init_weights_on_load = True @@ -862,9 +986,10 @@ class ModelCapabilities: """Declared parallelism capabilities for this model class.""" supports_tp: bool = False - supports_cp: bool = False + supports_cp: bool = True supports_pp: bool = False supports_ep: bool = True + supports_thd: bool = True @classmethod def from_config( @@ -928,6 +1053,33 @@ def update_moe_gate_bias(self) -> None: if isinstance(layer.mlp, MoE) and layer.mlp.gate.bias_update_factor > 0: layer.mlp.gate.update_bias() + def prepare_model_inputs_for_cp(self, batch: dict[str, Any], *, num_chunks: int = 1) -> dict[str, Any]: + """Select contiguous block-diagonal sharding for packed Laguna inputs. + + Args: + batch: Full packed batch with token tensors shaped [batch, sequence] + and THD ``seq_lens`` metadata. + num_chunks: Number of pipeline chunks; Laguna currently supports one. + + Returns: + A mapping containing the model-owned context-parallel sharder. + """ + if num_chunks != 1: + raise ValueError("Laguna packed context parallelism does not support pipeline microbatch chunking.") + if batch.get("qkv_format") != "thd": + raise ValueError("Laguna context parallelism requires packed THD inputs.") + if self.backend.attn != "sdpa": + raise ValueError("Laguna packed context parallelism requires model.backend.attn='sdpa'.") + + from nemo_automodel.components.distributed.blockdiag_cp import make_cp_blockdiag_batch_and_ctx + + return { + "cp_sharder": ContextParallelSharder( + shard_batch=partial(make_cp_blockdiag_batch_and_ctx, shard_primary=True), + local_token_global_indices=contiguous_local_indices, + ) + } + def forward( self, input_ids: torch.Tensor | None = None, @@ -945,9 +1097,12 @@ def forward( """Run the Laguna causal language model. Args: - input_ids: Optional token IDs of shape [batch, sequence]. - inputs_embeds: Optional embeddings of shape [batch, sequence, hidden]. - position_ids: Optional position IDs of shape [batch, sequence]. + input_ids: Optional token IDs of shape [batch, sequence] for BSHD or + [1, total_tokens] for packed THD input preparation. + inputs_embeds: Optional embeddings of shape [batch, sequence, hidden] + for BSHD or [1, total_tokens, hidden] for THD preparation. + position_ids: Optional position IDs of shape [batch, sequence] or + [1, total_tokens] for THD. attention_mask: Optional 2D bool/int mask of shape [batch, sequence], a 4D additive mask, or a mapping with per-attention-type masks keyed by "full_attention" and "sliding_attention". @@ -958,7 +1113,8 @@ def forward( logits_to_keep: If 0, compute logits for all sequence positions. If an int or tensor, compute logits only for the selected trailing positions. output_hidden_states: When true, include final hidden states in the output. - **kwargs: Additional attention backend arguments. + **kwargs: Additional attention backend arguments. THD requires + ``qkv_format='thd'`` and cumulative document lengths. Returns: Causal LM output with logits and optional hidden states. @@ -970,6 +1126,19 @@ def forward( if output_hidden_states is not None else getattr(self.config, "output_hidden_states", False) ) + is_thd = kwargs.get("qkv_format") == "thd" + if is_thd: + if position_ids is None: + raise ValueError("Packed THD input requires position_ids.") + input_ids, position_ids, padding_mask, kwargs = squeeze_input_for_thd( + input_ids, + position_ids, + padding_mask, + kwargs, + ) + if inputs_embeds is not None and inputs_embeds.ndim > 2: + inputs_embeds = inputs_embeds.squeeze(0) + attention_mask = None hidden = self.model( input_ids=input_ids, inputs_embeds=inputs_embeds, @@ -982,6 +1151,7 @@ def forward( self.lm_head, hidden, logits_to_keep, + is_thd=is_thd, output_hidden_states=output_hidden_states, ) diff --git a/tests/unit_tests/distributed/test_blockdiag_cp.py b/tests/unit_tests/distributed/test_blockdiag_cp.py index ac1b08be23..31191bc256 100644 --- a/tests/unit_tests/distributed/test_blockdiag_cp.py +++ b/tests/unit_tests/distributed/test_blockdiag_cp.py @@ -52,7 +52,7 @@ def apply(x, group, seq_dim): return x -def _run_blockdiag_cp(Q, K, V, doc_ids, world, enable_gqa=False): +def _run_blockdiag_cp(Q, K, V, doc_ids, world, enable_gqa=False, window_size=None): """Simulate ``world`` CP ranks in-process; returns the concatenated local outputs. Args: @@ -62,6 +62,7 @@ def _run_blockdiag_cp(Q, K, V, doc_ids, world, enable_gqa=False): doc_ids: Per-position document ids ``[B, S]`` (0 == padding). world: Number of simulated CP ranks. enable_gqa: Forwarded to the SDPA under test. + window_size: Optional sliding-attention window. Returns: Concatenated per-rank outputs ``[B, Hq, S, D]``. @@ -83,7 +84,15 @@ def _run_blockdiag_cp(Q, K, V, doc_ids, world, enable_gqa=False): token = bd_state._CP_BLOCKDIAG_STATE.set(state) try: q_local = Q[:, :, r * L : (r + 1) * L, :] - outs.append(bd_runtime.cp_blockdiag_sdpa(q_local, K, V, enable_gqa=enable_gqa)) + outs.append( + bd_runtime.cp_blockdiag_sdpa( + q_local, + K, + V, + enable_gqa=enable_gqa, + window_size=window_size, + ) + ) finally: bd_state._CP_BLOCKDIAG_STATE.reset(token) finally: @@ -169,6 +178,29 @@ def test_blockdiag_mask_expected_matrix(): assert torch.equal(got, expected) +def test_blockdiag_mask_applies_sliding_window_within_each_document(): + doc_ids = _doc_ids() + got = bd_kernels._cp_blockdiag_mask(doc_ids, 0, 8, 8, 1, window_size=(1, 0)) + expected = torch.eye(8, dtype=torch.bool) + expected[1, 0] = True + expected[2, 1] = True + expected[4, 3] = True + expected[5, 4] = True + + assert torch.equal(got, expected.view(1, 1, 8, 8)) + + +def test_blockdiag_doc_ids_expand_thd_lengths_and_preserve_padding(): + batch = { + "seq_lens": torch.tensor([[3, 2, -1000]]), + "seq_lens_padded": torch.tensor([[4, 2, -1000]]), + } + + got = bd_batch._cp_blockdiag_doc_ids(batch, seq_len=8, device=torch.device("cpu"), batch_size=1) + + assert torch.equal(got, torch.tensor([[1, 1, 1, 0, 2, 2, 0, 0]])) + + @pytest.mark.parametrize("world", [2, 4]) def test_blockdiag_sdpa_parity_vs_full_attention(world): """cp=world block-diagonal attention == cp=1 full attention on identical inputs.""" @@ -205,6 +237,23 @@ def test_blockdiag_sdpa_parity_gqa(world): assert max_diff < 1e-5, f"GQA world={world} CP-vs-full SDPA max_diff={max_diff}" +@pytest.mark.parametrize("world", [2, 4]) +def test_blockdiag_sdpa_sliding_window_parity(world): + torch.manual_seed(2) + B, H, S, D = 1, 4, 8, 16 + query = torch.randn(B, H, S, D, dtype=torch.float32) + key = torch.randn(B, H, S, D, dtype=torch.float32) + value = torch.randn(B, H, S, D, dtype=torch.float32) + doc_ids = _doc_ids() + window_size = (1, 0) + + cp_out = _run_blockdiag_cp(query, key, value, doc_ids, world=world, window_size=window_size) + mask = bd_kernels._cp_blockdiag_mask(doc_ids, 0, S, S, B, window_size=window_size) + reference = bd_runtime._ORIGINAL_SDPA(query, key, value, attn_mask=mask, is_causal=False) + + torch.testing.assert_close(cp_out, reference, atol=1e-6, rtol=1e-6) + + def test_blockdiag_sdpa_noop_without_state(): """With no CP state set, cp_blockdiag_sdpa is a plain pass-through to stock SDPA.""" torch.manual_seed(2) diff --git a/tests/unit_tests/models/laguna/test_model.py b/tests/unit_tests/models/laguna/test_model.py index b95769aa01..686249a738 100644 --- a/tests/unit_tests/models/laguna/test_model.py +++ b/tests/unit_tests/models/laguna/test_model.py @@ -14,8 +14,14 @@ import pytest import torch +import torch.nn.functional as F +from torch import nn +from nemo_automodel.components.distributed.blockdiag_cp import exchange as blockdiag_exchange +from nemo_automodel.components.distributed.blockdiag_cp import state as blockdiag_state +from nemo_automodel.components.distributed.blockdiag_cp.state import BlockdiagCpModelState from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.laguna import model as laguna_model_module from nemo_automodel.components.models.laguna.config import LagunaConfig from nemo_automodel.components.models.laguna.model import LagunaForCausalLM from nemo_automodel.components.moe.layers import MoE @@ -33,6 +39,75 @@ def _backend() -> BackendConfig: ) +def _te_backend() -> BackendConfig: + return BackendConfig( + attn="te", + linear="torch", + rms_norm="torch_fp32", + experts="torch", + dispatcher="torch", + enable_hf_state_dict_adapter=True, + ) + + +def _sdpa_backend() -> BackendConfig: + backend = _backend() + backend.attn = "sdpa" + return backend + + +class _IdentityGather: + @staticmethod + def apply(tensor, group, seq_dim): + del group, seq_dim + return tensor + + +class _ReferencePackedAttention(nn.Module): + """Independent CPU reference for TE's causal variable-length THD attention.""" + + def __init__(self, scale: float, attention_dropout: float) -> None: + super().__init__() + self.scale = scale + self.attention_dropout = attention_dropout + + def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> torch.Tensor: + cu_seqlens = kwargs["cu_seqlens_q"].tolist() + left_window, right_window = kwargs.get("window_size", (-1, 0)) + outputs = [] + for start, end in zip(cu_seqlens, cu_seqlens[1:]): + query_doc = query[start:end].transpose(0, 1).unsqueeze(0) + key_doc = key[start:end].transpose(0, 1).unsqueeze(0) + value_doc = value[start:end].transpose(0, 1).unsqueeze(0) + attention_mask = None + is_causal = True + if left_window >= 0: + positions = torch.arange(end - start) + query_positions = positions[:, None] + key_positions = positions[None, :] + attention_mask = (key_positions >= query_positions - left_window) & ( + key_positions <= query_positions + right_window + ) + is_causal = False + output = F.scaled_dot_product_attention( + query_doc, + key_doc, + value_doc, + attn_mask=attention_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + enable_gqa=True, + scale=self.scale, + ) + outputs.append(output.squeeze(0).transpose(0, 1)) + return torch.cat(outputs) + + +def _reference_te_factory(**kwargs): + attention = _ReferencePackedAttention(kwargs["softmax_scale"], kwargs["attention_dropout"]) + return attention, attention.__call__ + + def _tiny_config() -> LagunaConfig: cfg = LagunaConfig( vocab_size=32, @@ -157,3 +232,123 @@ def test_laguna_attention_uses_per_layer_head_counts_and_per_head_gate(): assert layer0_attn.g_proj.weight.shape == (2, 16) assert layer1_attn.q_proj.weight.shape == (16, 16) assert layer1_attn.g_proj.weight.shape == (4, 16) + + +def test_laguna_packed_thd_matches_per_document_logits_and_gradients(monkeypatch): + """THD must preserve Laguna document boundaries, RoPE, gating, and backward.""" + monkeypatch.setattr(laguna_model_module, "initialize_attn_module_and_func", _reference_te_factory) + torch.manual_seed(1234) + reference_model = LagunaForCausalLM(_dense_tiny_config(), backend=_te_backend()).to(torch.float32).train() + packed_model = LagunaForCausalLM(_dense_tiny_config(), backend=_te_backend()).to(torch.float32).train() + packed_model.load_state_dict(reference_model.state_dict()) + + input_ids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + position_ids = torch.tensor([0, 1, 2, 0, 1, 2, 3, 4]) + cu_seqlens = torch.tensor([0, 3, 8], dtype=torch.int32) + reference_logits = torch.cat( + [ + reference_model( + input_ids[:3].unsqueeze(0), + position_ids=position_ids[:3].unsqueeze(0), + ).logits, + reference_model( + input_ids[3:].unsqueeze(0), + position_ids=position_ids[3:].unsqueeze(0), + ).logits, + ], + dim=1, + ) + reference_logits.square().sum().backward() + + packed_logits = packed_model( + input_ids.unsqueeze(0), + position_ids=position_ids.unsqueeze(0), + qkv_format="thd", + cu_seqlens=cu_seqlens.unsqueeze(0), + max_seqlen=torch.tensor([5], dtype=torch.int32), + ).logits + packed_logits.square().sum().backward() + + torch.testing.assert_close(packed_logits, reference_logits, atol=1e-5, rtol=1e-5) + reference_params = dict(reference_model.named_parameters()) + for name, packed_param in packed_model.named_parameters(): + reference_grad = reference_params[name].grad + assert reference_grad is not None, name + assert packed_param.grad is not None, name + torch.testing.assert_close(packed_param.grad, reference_grad, atol=2e-5, rtol=2e-4) + + capabilities = packed_model.ModelCapabilities() + assert capabilities.supports_cp is True + assert capabilities.supports_thd is True + + +def test_laguna_blockdiag_thd_matches_per_document_sliding_attention(monkeypatch): + """The production THD+CP dispatch must preserve Laguna's sliding window.""" + monkeypatch.setattr(blockdiag_exchange, "_AllGatherSeqDiff", _IdentityGather) + torch.manual_seed(4321) + reference_model = LagunaForCausalLM(_dense_tiny_config(), backend=_sdpa_backend()).to(torch.float32).train() + packed_model = LagunaForCausalLM(_dense_tiny_config(), backend=_sdpa_backend()).to(torch.float32).train() + packed_model.load_state_dict(reference_model.state_dict()) + + input_ids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + position_ids = torch.tensor([0, 1, 2, 0, 1, 2, 3, 4]) + reference_logits = torch.cat( + [ + reference_model(input_ids[:3].unsqueeze(0), position_ids=position_ids[:3].unsqueeze(0)).logits, + reference_model(input_ids[3:].unsqueeze(0), position_ids=position_ids[3:].unsqueeze(0)).logits, + ], + dim=1, + ) + reference_logits.square().sum().backward() + + step_state = { + "group": None, + "doc_ids": torch.tensor([[1, 1, 1, 2, 2, 2, 2, 2]]), + "row_offset": 0, + "seq_dim": 2, + "attn_backend": "dense", + "kv_exchange": "allgather", + "model_state": BlockdiagCpModelState( + group=None, + packed_cu_seqlens=torch.tensor([0, 3, 8]), + packed_cu_seqlens_cpu=torch.tensor([0, 3, 8]), + ), + } + token = blockdiag_state._CP_BLOCKDIAG_STATE.set(step_state) + try: + packed_logits = packed_model( + input_ids.unsqueeze(0), + position_ids=position_ids.unsqueeze(0), + qkv_format="thd", + ).logits + packed_logits.square().sum().backward() + finally: + blockdiag_state._CP_BLOCKDIAG_STATE.reset(token) + + torch.testing.assert_close(packed_logits, reference_logits, atol=1e-5, rtol=1e-5) + reference_params = dict(reference_model.named_parameters()) + for name, packed_param in packed_model.named_parameters(): + torch.testing.assert_close(packed_param.grad, reference_params[name].grad, atol=2e-5, rtol=2e-4) + + +def test_laguna_reports_sdpa_cp_and_packing_support(): + from nemo_automodel._transformers.capabilities import ModelSupports + + model = LagunaForCausalLM(_dense_tiny_config(), backend=_sdpa_backend()) + supports = ModelSupports(model, None) + + assert supports.supports_cp is True + assert supports.supports_sequence_packing is True + + +def test_laguna_packed_thd_rejects_non_te_attention(): + model = LagunaForCausalLM(_dense_tiny_config(), backend=_backend()).eval() + + with pytest.raises(ValueError, match="requires backend.attn='te'"): + model( + torch.tensor([[1, 2, 3, 4]]), + position_ids=torch.tensor([[0, 1, 0, 1]]), + qkv_format="thd", + cu_seqlens=torch.tensor([[0, 2, 4]], dtype=torch.int32), + max_seqlen=torch.tensor([2], dtype=torch.int32), + ) From 10e70902c1995fdcfcc9634137bdf7e59b11357d Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis Date: Sun, 23 Aug 2026 22:31:34 -0700 Subject: [PATCH 2/3] fix(laguna): preserve native TE THD input sharding Signed-off-by: Alexandros Koumparoulis --- nemo_automodel/components/models/laguna/model.py | 10 +++++++--- tests/unit_tests/models/laguna/test_model.py | 9 +++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/nemo_automodel/components/models/laguna/model.py b/nemo_automodel/components/models/laguna/model.py index 73f5ee44b2..e057d16beb 100644 --- a/nemo_automodel/components/models/laguna/model.py +++ b/nemo_automodel/components/models/laguna/model.py @@ -1054,7 +1054,7 @@ def update_moe_gate_bias(self) -> None: layer.mlp.gate.update_bias() def prepare_model_inputs_for_cp(self, batch: dict[str, Any], *, num_chunks: int = 1) -> dict[str, Any]: - """Select contiguous block-diagonal sharding for packed Laguna inputs. + """Select input preparation for packed Laguna inputs. Args: batch: Full packed batch with token tensors shaped [batch, sequence] @@ -1062,14 +1062,18 @@ def prepare_model_inputs_for_cp(self, batch: dict[str, Any], *, num_chunks: int num_chunks: Number of pipeline chunks; Laguna currently supports one. Returns: - A mapping containing the model-owned context-parallel sharder. + A mapping containing the model-owned context-parallel sharder for + SDPA. TE returns an empty mapping so the framework selects its native + THD sharder, including when context parallelism is disabled. """ if num_chunks != 1: raise ValueError("Laguna packed context parallelism does not support pipeline microbatch chunking.") if batch.get("qkv_format") != "thd": raise ValueError("Laguna context parallelism requires packed THD inputs.") + if self.backend.attn == "te": + return {} if self.backend.attn != "sdpa": - raise ValueError("Laguna packed context parallelism requires model.backend.attn='sdpa'.") + raise ValueError("Laguna packed THD input preparation requires model.backend.attn='te' or 'sdpa'.") from nemo_automodel.components.distributed.blockdiag_cp import make_cp_blockdiag_batch_and_ctx diff --git a/tests/unit_tests/models/laguna/test_model.py b/tests/unit_tests/models/laguna/test_model.py index 686249a738..fd1d70f322 100644 --- a/tests/unit_tests/models/laguna/test_model.py +++ b/tests/unit_tests/models/laguna/test_model.py @@ -341,6 +341,15 @@ def test_laguna_reports_sdpa_cp_and_packing_support(): assert supports.supports_sequence_packing is True +def test_laguna_te_thd_uses_framework_input_sharder(): + model = LagunaForCausalLM(_dense_tiny_config(), backend=_sdpa_backend()) + model.backend.attn = "te" + + prepared = model.prepare_model_inputs_for_cp({"qkv_format": "thd"}) + + assert prepared == {} + + def test_laguna_packed_thd_rejects_non_te_attention(): model = LagunaForCausalLM(_dense_tiny_config(), backend=_backend()).eval() From 0ee3a52b69d9adefa79929cea0a04dded1856086 Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:35:05 -0700 Subject: [PATCH 3/3] Apply suggestions from code review Co-authored-by: jgerh <163925524+jgerh@users.noreply.github.com> --- docs/model-coverage/llm/poolside/laguna.mdx | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/model-coverage/llm/poolside/laguna.mdx b/docs/model-coverage/llm/poolside/laguna.mdx index c5fd539129..00fe2d6dc4 100644 --- a/docs/model-coverage/llm/poolside/laguna.mdx +++ b/docs/model-coverage/llm/poolside/laguna.mdx @@ -1,9 +1,9 @@ --- title: "Laguna" -description: "" +description: "Laguna Mixture-of-Experts language models in NeMo AutoModel, including Laguna S 2.1, Laguna XS 2.1, expert parallelism, and THD sequence packing." slug: model-coverage/large-language-models/poolside/laguna --- -[Laguna](https://huggingface.co/poolside) is Poolside's hybrid-attention MoE language model family. The Automodel implementation supports full SFT with expert parallelism for Laguna S 2.1 and Laguna XS 2.1, including native THD sequence packing with context parallelism. +[Laguna](https://huggingface.co/poolside) is Poolside's hybrid-attention Mixture-of-Experts (MoE) language model family. NeMo AutoModel supports full SFT with expert parallelism for Laguna S 2.1 and Laguna XS 2.1, including native THD sequence packing with context parallelism. @@ -28,9 +28,9 @@ slug: model-coverage/large-language-models/poolside/laguna | Recipe | Description | |---|---| -| [laguna_s_2p1_hellaswag_ep16.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_s_2p1_hellaswag_ep16.yaml) | SFT — Laguna S 2.1 on HellaSwag with EP16 | -| [laguna_xs_2p1_hellaswag_ep8.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml) | SFT — Laguna XS 2.1 on HellaSwag with EP8 and 1K NEAT sequence packing | -| [laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml) | SFT — Laguna XS 2.1 on HellaSwag with EP8, CP2, and 1K THD sequence packing | +| [laguna_s_2p1_hellaswag_ep16.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_s_2p1_hellaswag_ep16.yaml) | SFT, Laguna S 2.1 on HellaSwag with EP16 | +| [laguna_xs_2p1_hellaswag_ep8.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8.yaml) | SFT, Laguna XS 2.1 on HellaSwag with EP8 and 1K NEAT sequence packing | +| [laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/laguna/laguna_xs_2p1_hellaswag_ep8_cp2_thd.yaml) | SFT, Laguna XS 2.1 on HellaSwag with EP8, CP2, and 1K THD sequence packing | ## Run the Recipe