feat(laguna): support packed THD context parallelism

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.
 
 <Info>
 
@@ -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,15 +30,19 @@ 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
 
 ```bash
 # 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`.

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
 

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"

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.

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]

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,