refactor(cp): sharder ownership

nemo_automodel/_transformers/capabilities.py

@@ -98,7 +98,7 @@ def _uses_magi_attention(model: "nn.Module") -> bool:
     """True when the model uses the MagiAttention (FFA / context-parallel) backend.
 
     MagiAttention implements context parallelism via its own load-balancing
-    dispatch (see ``components/distributed/magi_attn_utils.py``), so it supports CP.
+    dispatch (see ``components/distributed/context_parallel/magi.py``), so it supports CP.
     """
     backend = getattr(model, "backend", None)
     return getattr(backend, "attn", None) == "magi"
@@ -336,7 +336,7 @@ def supports_cp_with_sequence_packing(self) -> bool:
 
         MagiAttention dispatches the packed sequence across the CP group with its
         own load-balancing solver and a per-document varlen mask, so it supports
-        CP + packing (see ``magi_attn_utils.magi_prepare_packed_cp``). Models
+        CP + packing (see ``context_parallel.magi.magi_prepare_packed_cp``). Models
         with native THD support own their packed CP path in TileLang attention."""
         model = self._model
         if not self.supports_sequence_packing:

nemo_automodel/_transformers/infrastructure.py

@@ -748,7 +748,7 @@ def apply_model_infrastructure(
     # is not excluded by the _uses_te_attention check, so gate on ep_size: only
     # dense (non-MoE) models need this pass.
     if mesh.cp_size > 1 and mesh.ep_size <= 1 and not _uses_te_attention(model):
-        from nemo_automodel.components.distributed.cp_utils import (
+        from nemo_automodel.components.distributed.context_parallel.utils import (
             attach_context_parallel_hooks,
             attach_cp_sdpa_hooks,
         )

nemo_automodel/components/attention/utils.py

@@ -108,7 +108,7 @@ def attn_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **call_kwargs:
             )
         # requires magi_attention; the guards above are exercised on CPU but the
         # kernel build is not, so exclude it from coverage.
-        from nemo_automodel.components.distributed.magi_attn_utils import (  # pragma: no cover - requires magi_attention
+        from nemo_automodel.components.distributed.context_parallel.magi import (  # pragma: no cover - requires magi_attention
             make_magi_attn_func,
         )
 

nemo_automodel/components/datasets/loader.py

@@ -589,6 +589,13 @@ def dataset_builds_on_all_ranks(self) -> bool:
         """Whether dataset construction must bypass rank-zero-first ordering."""
         return isinstance(self.dataset_config, AllRanksDatasetConfig)
 
+    @property
+    def emits_thd(self) -> bool:
+        """Whether this configuration produces THD-formatted batches."""
+        from nemo_automodel.components.datasets.utils import packed_sequence_thd_collater
+
+        return isinstance(self.packing, ThdPackingConfig) or self.collate_fn is packed_sequence_thd_collater
+
     def _build_dataset(
         self,
         *,

nemo_automodel/components/distributed/context_parallel/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+"""Context-parallel batch sharding."""
+
+from nemo_automodel.components.distributed.context_parallel.sharder import ContextParallelSharder
+
+__all__ = ["ContextParallelSharder"]

nemo_automodel/components/distributed/magi_attn_utils.py → nemo_automodel/components/distributed/context_parallel/magi.py

@@ -50,7 +50,7 @@
 import torch
 import torch.distributed as dist
 
-from nemo_automodel.components.distributed.cp_utils import _make_cp_batch_and_ctx, make_cp_batch_for_te
+from nemo_automodel.components.distributed.context_parallel.utils import _make_cp_batch_and_ctx, make_cp_batch_for_te
 
 logger = logging.getLogger(__name__)
 
@@ -753,7 +753,7 @@ def prepare_llm_batch(
         context). ``local_indices`` is the global stream position of every
         local token on the paths that dispatch the sequence (magi's
         ``get_position_ids``), None otherwise; the framework installs it on
-        the magi ContextParallelismSharder for the token-tensor verbs.
+        the magi ContextParallelSharder for the token-tensor verbs.
         """
         # cp=1 prefix-tree mask: the datasets layer cannot import this module (component
         # independence), so the collate attaches the tree structure and the spec is built
@@ -843,7 +843,7 @@ def make_cp_batch(
             return_local_indices: Also return the local-token global index map
                 from the dispatch that just ran (None on paths that do not
                 dispatch, e.g. cp=1 THD conversion and the VLM domain). Used
-                by the magi ContextParallelismSharder's token verbs.
+                by the magi ContextParallelSharder's token verbs.
 
         Returns:
             The dispatched (magi-sharded) batch, or ``(batch, local_indices)``

nemo_automodel/components/distributed/cp_sharder.py → nemo_automodel/components/distributed/context_parallel/sharder.py

@@ -14,13 +14,13 @@
 
 """Context-parallel batch-sharding contract.
 
-Every CP backend is a :class:`ContextParallelismSharder`. A model that owns its CP batch
+Every CP backend is a :class:`ContextParallelSharder`. A model that owns its CP batch
 sharding and attention transport returns one from
 ``prepare_model_inputs_for_cp`` under the ``"cp_sharder"`` batch key; the
 framework constructs its own for the remaining backends (torch
 ``context_parallel`` round-robin, TE/THD, MagiAttention) so
-the CP dispatch (``cp_utils.prepare_cp_forward``) reduces to resolving a sharder and calling
-``shard_batch``. This replaces the retired private batch keys
+constructing :class:`ContextParallelSharder` resolves and configures the backend;
+callers then invoke ``sharder.shard(batch)``. This replaces the retired private batch keys
 (``_cp_make_batch_fn``, ``_cp_metadata_seq_dims``, ``_cp_metadata_pad_values``,
 ``_cp_full_logits_grad_touch``).
 
@@ -44,7 +44,7 @@
 shared contiguous-shard batch prep used by models whose CP ranks own contiguous
 sequence slices (Gemma4, DeepSeek V4), and the torch ``context_parallel``
 round-robin load-balanced prep with its index map. The TE/THD and magi preps
-live with their dependencies (``cp_utils``, ``magi_attn_utils``); the
+live with their dependencies (``context_parallel.utils``, ``context_parallel.magi``); the
 dispatcher wraps them into sharders at resolution time.
 """
 
@@ -57,6 +57,7 @@
 
 import torch
 import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
 
 
 def _cp_rank(cp_mesh) -> int:
@@ -236,38 +237,141 @@ class ShardLayout:
     input_token_stream_positions: torch.Tensor | None = None
 
 
-@dataclass
-class ContextParallelismSharder:
+class ContextParallelSharder:
     """CP backend description: how a batch is sharded and where local tokens live.
 
     Attributes:
         shard_batch: ``(cp_mesh, tp_mesh, batch, *, loss_mask=None,
             padding_token_id=0) -> (ctx_factory, batch, ShardLayout | None)``.
             Pads and shards the batch, installs any backend-owned attention
-            transport, and reports the shard layout it computed; the dispatch
+            transport, and reports the shard layout it computed; :meth:`shard`
             stores it as ``shard_layout`` for the token verbs.
         local_token_global_indices: ``(cp_mesh, padded_seq_len, device) ->
             LongTensor`` with the global position of each local token —
             closed-form for contiguous/round-robin layouts; None for
             data-dependent layouts, whose partition arrives with
             ``shard_layout`` (their token verbs raise before the first shard).
         shard_layout: The :class:`ShardLayout` of the last ``shard_batch``, set
-            by the dispatch. Sharders are built per resolution/hook call, so
+            by :meth:`shard`. Sharders are built per resolution/hook call, so
             the layout never leaks across steps.
     """
 
     shard_batch: Callable[..., tuple[Callable, dict[str, Any], "ShardLayout | None"]]
     local_token_global_indices: Callable[..., torch.Tensor] | None
-    shard_layout: "ShardLayout | None" = None
+    shard_layout: "ShardLayout | None"
+    _cp_mesh: Any
+    _tp_mesh: Any
+    _loss_mask: torch.Tensor | None
+    _padding_token_id: int
 
-    def _indices(self, cp_mesh, padded_seq_len: int, device) -> torch.Tensor:
+    def __init__(
+        self,
+        model: torch.nn.Module | None = None,
+        device_mesh: DeviceMesh | None = None,
+        batch: dict[str, Any] | None = None,
+        *,
+        shard_batch: Callable[..., tuple[Callable, dict[str, Any], "ShardLayout | None"]] | None = None,
+        local_token_global_indices: Callable[..., torch.Tensor] | None = None,
+        shard_layout: "ShardLayout | None" = None,
+        padding_token_id: int = 0,
+        num_chunks: int = 1,
+        loss_mask: torch.Tensor | None = None,
+        invoke_pre_embed: bool = True,
+        extra_seq_buffers: dict[str, int] | None = None,
+    ) -> None:
+        """Construct a strategy sharder or resolve one for a forward.
+
+        Args:
+            model: Model whose attention backend and CP preparation hook select
+                the sharding strategy, or None for the generic strategy. Omit
+                when constructing directly from ``shard_batch``.
+            device_mesh: Device mesh containing optional ``cp`` and ``tp`` axes.
+                Direct strategy construction uses it to configure the sharder;
+                omit it only when returning an unresolved model-owned strategy.
+            batch: Mutable input mapping required when resolving a strategy.
+                Token tensors normally have shape
+                [batch, sequence, ...]; THD source batches declare
+                ``qkv_format="thd"`` and are flattened during :meth:`shard`.
+                Model hook metadata is merged into this mapping in place.
+            shard_batch: Optional backend callback for direct strategy
+                construction. It accepts token tensors with backend-defined
+                layouts and returns the sharded batch plus its layout.
+            local_token_global_indices: Optional callback returning a tensor of
+                shape [local_tokens] with each local token's global position.
+            shard_layout: Optional captured layout for direct strategy
+                construction. Tensor fields use the layouts documented by
+                :class:`ShardLayout`.
+            padding_token_id: Value used to pad ``input_ids`` on the sequence axis.
+            num_chunks: Number of THD chunks created during sharding.
+            loss_mask: Optional tensor of shape [batch, sequence] sharded with
+                the batch.
+            invoke_pre_embed: Whether to invoke a model-owned CP preparation hook.
+            extra_seq_buffers: Additional batch keys mapped to their sequence axes.
+        """
+        if shard_batch is not None:
+            has_resolution_args = (
+                model is not None
+                or batch is not None
+                or num_chunks != 1
+                or not invoke_pre_embed
+                or extra_seq_buffers is not None
+            )
+            if has_resolution_args:
+                raise TypeError("shard_batch is mutually exclusive with model, batch, and strategy-resolution options")
+
+            self.shard_batch = shard_batch
+            self.local_token_global_indices = local_token_global_indices
+            self.shard_layout = shard_layout
+            mesh_dim_names = getattr(device_mesh, "mesh_dim_names", ())
+            self._cp_mesh = device_mesh["cp"] if "cp" in mesh_dim_names else None
+            self._tp_mesh = device_mesh["tp"] if "tp" in mesh_dim_names else None
+            self._loss_mask = loss_mask
+            self._padding_token_id = padding_token_id
+            return
+        if local_token_global_indices is not None or shard_layout is not None:
+            raise TypeError("local_token_global_indices and shard_layout require shard_batch")
+        if batch is None:
+            raise TypeError("batch is required when shard_batch is not provided")
+
+        from nemo_automodel.components.distributed.context_parallel.utils import _prepare_cp_sharder
+
+        resolved = _prepare_cp_sharder(
+            model,
+            device_mesh,
+            batch,
+            padding_token_id=padding_token_id,
+            num_chunks=num_chunks,
+            loss_mask=loss_mask,
+            invoke_pre_embed=invoke_pre_embed,
+            extra_seq_buffers=extra_seq_buffers,
+        )
+        self.shard_batch = resolved.shard_batch
+        self.local_token_global_indices = resolved.local_token_global_indices
+        self.shard_layout = resolved.shard_layout
+        self._cp_mesh = resolved._cp_mesh
+        self._tp_mesh = resolved._tp_mesh
+        self._loss_mask = resolved._loss_mask
+        self._padding_token_id = resolved._padding_token_id
+
+    def shard(self, batch: dict[str, Any]) -> tuple[Callable, dict[str, Any]]:
+        """Shard a batch and retain its layout for token-aligned tensors."""
+        ctx, batch, self.shard_layout = self.shard_batch(
+            self._cp_mesh,
+            self._tp_mesh,
+            batch,
+            loss_mask=self._loss_mask,
+            padding_token_id=self._padding_token_id,
+        )
+        return ctx, batch
+
+    def _indices(self, padded_seq_len: int, device) -> torch.Tensor:
         layout = self.shard_layout or _NO_SHARD_LAYOUT
         captured = layout.local_token_global_indices
         if captured is not None:
             # Data-dependent layout: use the partition the shard reported, and
             # validate the requested length against it so a mismatched tensor
             # cannot be silently mis-sharded.
-            cp_size = cp_mesh.size() if cp_mesh is not None else 1
+            cp_size = self._cp_mesh.size() if self._cp_mesh is not None else 1
             expected = captured.numel() * cp_size
             if padded_seq_len != expected:
                 raise ValueError(
@@ -278,14 +382,14 @@ def _indices(self, cp_mesh, padded_seq_len: int, device) -> torch.Tensor:
             return captured.reshape(-1).to(device=device, dtype=torch.long)
         if self.local_token_global_indices is None:
             raise NotImplementedError(
-                "This ContextParallelismSharder has a data-dependent token layout; its index map "
+                "This ContextParallelSharder has a data-dependent token layout; its index map "
                 "arrives with the shard layout — token-tensor shard/gather are unavailable before "
                 "the first shard."
             )
-        return self.local_token_global_indices(cp_mesh, padded_seq_len, device)
+        return self.local_token_global_indices(self._cp_mesh, padded_seq_len, device)
 
     def shard_token_tensor(
-        self, cp_mesh, tensor: torch.Tensor, seq_dim: int = 1, fill: float | int | None = None
+        self, tensor: torch.Tensor, seq_dim: int = 1, fill: float | int | None = None
     ) -> torch.Tensor:
         """Shard a full-length token-aligned tensor exactly like the model inputs.
 
@@ -330,16 +434,15 @@ def shard_token_tensor(
                 tensor = _pad_tensor_seq_dim_(tensor, seq_dim, layout.padded_seq_len - length, fill)
             else:
                 raise ValueError(
-                    f"This ContextParallelismSharder sharded a batch of padded_seq_len={layout.padded_seq_len} "
+                    f"This ContextParallelSharder sharded a batch of padded_seq_len={layout.padded_seq_len} "
                     f"(original_seq_len={layout.original_seq_len}), got a tensor of length {length} on dim {seq_dim}. "
                     "Pass the original-length tensor with an explicit `fill`, or pre-pad it yourself."
                 )
-        indices = self._indices(cp_mesh, tensor.shape[seq_dim], tensor.device)
+        indices = self._indices(tensor.shape[seq_dim], tensor.device)
         return shard_token_tensor_by_indices(tensor, indices, seq_dim=seq_dim)
 
     def gather_token_tensor(
         self,
-        cp_mesh,
         tensor: torch.Tensor,
         seq_dim: int = 1,
         trim: bool = False,
@@ -355,9 +458,9 @@ def gather_token_tensor(
         (nothing to trim to).
         """
         layout = self.shard_layout or _NO_SHARD_LAYOUT
-        padded_seq_len = tensor.shape[seq_dim] * (cp_mesh.size() if cp_mesh is not None else 1)
-        indices = self._indices(cp_mesh, padded_seq_len, tensor.device)
-        full = gather_token_tensor_by_indices(cp_mesh, tensor, indices, seq_dim=seq_dim)
+        padded_seq_len = tensor.shape[seq_dim] * (self._cp_mesh.size() if self._cp_mesh is not None else 1)
+        indices = self._indices(padded_seq_len, tensor.device)
+        full = gather_token_tensor_by_indices(self._cp_mesh, tensor, indices, seq_dim=seq_dim)
         if not trim:
             return full
         if layout.padded_seq_len is not None and full.shape[seq_dim] != layout.padded_seq_len:
@@ -377,7 +480,7 @@ def gather_token_tensor(
         if layout.original_seq_len is not None:
             return full.narrow(seq_dim, 0, layout.original_seq_len)
         raise NotImplementedError(
-            "This ContextParallelismSharder has no shard layout to trim to; "
+            "This ContextParallelSharder has no shard layout to trim to; "
             "gather with trim=False and restore the layout with the batch metadata "
             "(padding_mask / cu_seqlens)."
         )
@@ -718,7 +821,7 @@ def shard_batch_load_balanced(
 ):
     """Shard a batch with torch ``context_parallel`` round-robin load balancing.
 
-    ``ContextParallelismSharder.shard_batch`` implementation for the default framework-owned CP
+    ``ContextParallelSharder.shard_batch`` implementation for the default framework-owned CP
     path (layout ``"round_robin"``, indices from
     :func:`round_robin_local_indices`). Assumes an active CP mesh (size > 1).
     ``padding_token_id`` is accepted per the contract but unused: CP-pad slots
@@ -729,9 +832,8 @@ def shard_batch_load_balanced(
         ``(ctx_factory, batch, ShardLayout)`` where entering ``ctx_factory()``
         installs the SDPA-kernel + ``context_parallel`` context for the forward.
     """
-    # Call-time import: the torch-CP transport machinery stays in cp_utils
-    # (NeMo-RL imports it from there), and cp_utils imports this module.
-    from nemo_automodel.components.distributed.cp_utils import (  # noqa: PLC0415
+    # Call-time import avoids a cycle: utils imports this module's strategy helpers.
+    from nemo_automodel.components.distributed.context_parallel.utils import (  # noqa: PLC0415
         _shard_grad_buffer_for_cp,
         create_context_parallel_ctx,
         get_train_context,
@@ -865,7 +967,7 @@ def shard_batch_aux_only(
         The layout's ``padded_seq_len`` is what the model must pad its primary
         stream to before sharding.
     """
-    from nemo_automodel.components.distributed.cp_utils import (  # noqa: PLC0415
+    from nemo_automodel.components.distributed.context_parallel.utils import (  # noqa: PLC0415
         create_context_parallel_ctx,
         get_train_context,
     )