feat: HYBRID_SHARD draft training across colocated target islands

docs/basic_usage/colocated_training.md

@@ -55,6 +55,41 @@ Set `model.sglang_context_length` to at least `data.max_length + 7`; size
 `training.tp_size * training.batch_size` requests per island. Explicit values
 override the safe derived defaults.
 
+## HSDP across islands
+
+For multi-node targets, `training.fsdp_sharding: HYBRID_SHARD` uses the target
+TP group as the FSDP shard group and the target-DP group for replication:
+
+| Target shape | Suggested topology | Draft sharding |
+| --- | --- | --- |
+| Qwen3-8B on 8 H200 | 8 islands of TP1 | `SHARD_GRAD_OP` across all 8 ranks |
+| One-node target | one TP island per node | `HYBRID_SHARD` |
+| K3-class target on 4x8 B300 | 4 islands of TP8 | `HYBRID_SHARD` |
+
+For a TP8 target on four eight-GPU nodes:
+
+```yaml
+training:
+  batch_size: 1
+  tp_size: 8
+  fsdp_sharding: HYBRID_SHARD
+
+deployment:
+  mode: local_colocated
+  trainer:
+    nnodes: 4
+    nproc_per_node: 8
+    master_addr: trainer-0
+```
+
+The complete K3 starting recipe is
+[`kimi-k3-dspark-colocated.yaml`](../../examples/configs/online/colocated/kimi-k3-dspark-colocated.yaml).
+
+HSDP shards the draft inside each TP8 island and replicates corresponding
+shards across islands. This keeps parameter all-gathers node-local while only
+the replica synchronization crosses nodes. Loss and metric reductions remain
+WORLD-wide, and the gradient norm counts each replicated shard once.
+
 Prompt-cache preparation is coordinated: rank zero builds the tokenized Arrow
 cache first (node-local caches are then built once per remaining node) and all
 other ranks take the cache-hit path. The coordination collective runs inside

examples/configs/README.md

@@ -253,7 +253,7 @@ Common fields:
 | `training.total_steps` | `null` | Positive optimizer/loss schedule horizon; it does not itself stop an online stream. A finite online disaggregated run may omit both fields: the producer publishes the exact horizon derived from prepared prompts, epochs, DP size, batch size, and accumulation. |
 | `training.batch_size` | `1` | Per-rank microbatch size. P-EAGLE and USP require 1. |
 | `training.accumulation_steps` | `1` | Positive microbatches per optimizer update. |
-| `training.fsdp_sharding` | `SHARD_GRAD_OP` | Trainer FSDP mode: `SHARD_GRAD_OP`, `FULL_SHARD`, or `NO_SHARD`. |
+| `training.fsdp_sharding` | `SHARD_GRAD_OP` | Trainer FSDP mode: `SHARD_GRAD_OP`, `FULL_SHARD`, `HYBRID_SHARD`, or `NO_SHARD`. `HYBRID_SHARD` shards over the colocated target-TP group and replicates across target-DP islands. |
 | `training.learning_rate` | `1e-4` | Positive peak learning rate. |
 | `training.lr_scheduler` | `cosine` | Learning-rate schedule after warmup: `cosine` or `constant`. |
 | `training.warmup_ratio` | `0.015` | Fraction in `[0, 1]` used for scheduler warmup. |

examples/configs/online/colocated/README.md

@@ -8,6 +8,7 @@ and the trainer world size must be divisible by it.
 | Recipe | Target | Topology |
 | --- | --- | --- |
 | [`qwen3-8b-dspark-colocated.yaml`](qwen3-8b-dspark-colocated.yaml) | Qwen3-8B | 8 islands of TP1 on one 8xH200 node |
+| [`kimi-k3-dspark-colocated.yaml`](kimi-k3-dspark-colocated.yaml) | Kimi K3 | 4 islands of TP8 with `HYBRID_SHARD` on 4x8 B300 |
 
 See [Colocated online training](../../../../docs/basic_usage/colocated_training.md)
 for memory sizing and scaling guidance.

examples/configs/online/colocated/kimi-k3-dspark-colocated.yaml

@@ -0,0 +1,70 @@
+# Kimi K3 DSpark, online colocated — four TP8/HSDP islands on 4x8 B300.
+model:
+  target_model_path: /workspace/models/Kimi-K3-cdd2e49a
+  draft_model_config: configs/kimi-k3-dspark.json
+  draft_checkpoint_path: /workspace/k3_dspark/checkpoints/kimi-k3-dspark/epoch_0_step_0
+  target_backend: sglang
+  trust_remote_code: true
+  embedding_key: language_model.model.embed_tokens.weight
+  lm_head_key: language_model.lm_head.weight
+  mask_token_id: 163824
+  torch_dtype: bfloat16
+  sglang_attention_backend: trtllm_mla
+  sglang_mem_fraction_static: 0.76
+  sglang_disable_radix_cache: false
+  sglang_context_length: 66048
+  # A TP8 island captures one distinct rank-local sequence for every peer.
+  sglang_max_running_requests: 8
+  sglang_max_total_tokens: 528384
+  sglang_moe_runner_backend: marlin
+  sglang_mamba_radix_cache_strategy: extra_buffer
+  sglang_max_mamba_cache_size: 40
+
+data:
+  train_data_path: /workspace/k3_dspark/data/kimi-k3-agentic-regen.jsonl
+  max_length: 65536
+  chat_template: kimi-k3-thinking
+  cache_dir: /workspace/k3_dspark/cache
+  build_dataset_num_proc: 64
+  dataloader_num_workers: 0
+
+training:
+  strategy: dspark
+  num_epochs: 10
+  # 32 rank-local sequences x 4 microbatches = effective global batch 128.
+  batch_size: 1
+  accumulation_steps: 4
+  tp_size: 8
+  fsdp_sharding: HYBRID_SHARD
+  learning_rate: 0.000050959167111070076
+  lr_scheduler: constant
+  warmup_ratio: 0
+  max_grad_norm: 1
+  attention_backend: flex_attention
+  num_anchors: 512
+  loss_decay_gamma: 4.0
+  objective_chunk_blocks: 128
+  dspark_ce_loss_alpha: 0.1
+  dspark_l1_loss_alpha: 0.9
+  dspark_confidence_head_alpha: 1.0
+  save_interval: 8
+  log_interval: 10
+  dist_timeout: 30
+  seed: 42
+  prompt_seed: 1
+
+tracking:
+  report_to: wandb
+  wandb_project: specforge-dspark
+  wandb_name: kimi-k3-dspark-specforge-colocated-hsdp
+  wandb_dir: /workspace/k3_dspark/runs/kimi-k3-specforge/wandb
+
+run_id: kimi-k3-dspark-specforge-colocated
+output_dir: /workspace/k3_dspark/runs/kimi-k3-specforge/colocated-output
+
+deployment:
+  mode: local_colocated
+  trainer:
+    nnodes: 4
+    nproc_per_node: 8
+    master_addr: trainer-0

specforge/config/schema.py

@@ -527,7 +527,9 @@ class TrainingConfig(StrictConfigModel):
     total_steps: Optional[int] = Field(default=None, gt=0)
     batch_size: int = Field(default=1, gt=0)
     accumulation_steps: int = Field(default=1, gt=0)
-    fsdp_sharding: Literal["SHARD_GRAD_OP", "FULL_SHARD", "NO_SHARD"] = "SHARD_GRAD_OP"
+    fsdp_sharding: Literal[
+        "SHARD_GRAD_OP", "FULL_SHARD", "HYBRID_SHARD", "NO_SHARD"
+    ] = "SHARD_GRAD_OP"
     learning_rate: float = Field(default=1e-4, gt=0.0)
     lr_scheduler: Literal["cosine", "constant"] = "cosine"
     warmup_ratio: float = Field(default=0.015, ge=0.0, le=1.0)
@@ -961,6 +963,20 @@ def _validate_run_structure(self):
                     "model.sglang_ep_size must be no larger than and evenly "
                     "divide training.tp_size for colocated capture"
                 )
+        if self.training.fsdp_sharding == "HYBRID_SHARD":
+            trainer_world = (
+                self.deployment.trainer.nnodes * self.deployment.trainer.nproc_per_node
+            )
+            if deployment != "local_colocated" or mode != "online":
+                raise ValueError(
+                    "training.fsdp_sharding=HYBRID_SHARD currently requires "
+                    "online deployment.mode=local_colocated"
+                )
+            if self.training.tp_size <= 1 or trainer_world <= self.training.tp_size:
+                raise ValueError(
+                    "HYBRID_SHARD requires training.tp_size > 1 and at least "
+                    "two colocated target islands"
+                )
         if self.training.role == "producer" and self.training.resume_from is not None:
             raise ValueError("training.resume_from is valid only for a trainer role")
         if self.training.attention_backend == "usp":
@@ -1008,6 +1024,13 @@ def validate_world_size(self, world_size: int) -> None:
                 f"parallel size {sp_size} "
                 "(sp_ulysses_size * sp_ring_size)"
             )
+        if self.training.fsdp_sharding == "HYBRID_SHARD" and (
+            tp_size <= 1 or world_size <= tp_size
+        ):
+            raise ValueError(
+                "HYBRID_SHARD requires a target-TP shard group larger than one "
+                "and at least two target-DP replicas"
+            )
 
     @classmethod
     def from_file(cls, path: str) -> "Config":

specforge/training/backend.py

@@ -43,6 +43,13 @@ class ParallelConfig:
     sharding_strategy: str = "SHARD_GRAD_OP"
     param_dtype: torch.dtype = torch.bfloat16
     fsdp_process_group: Any = None
+    #: Group for loss/metric reductions; falls back to fsdp_process_group.
+    #: HSDP separates them because its FSDP group is a (shard, replica) pair.
+    reduction_process_group: Any = None
+    #: Group over which each parameter shard exists exactly once; HSDP
+    #: replicates shards across islands, so counting them WORLD-wide would
+    #: overstate the global gradient norm.
+    grad_norm_process_group: Any = None
     dp_group: Any = None
     draft_dp_group: Any = None
     tp_group: Any = None
@@ -111,14 +118,36 @@ def from_distributed(
                 "ParallelConfig.from_distributed: distributed handles unavailable: %s",
                 exc,
             )
+        fsdp_process_group: Any = dist.group.WORLD
+        grad_norm_process_group: Any = dist.group.WORLD
+        if sharding_strategy == "HYBRID_SHARD":
+            shard_group = handles.get("tp_group")
+            replica_group = handles.get("dp_group")
+            if shard_group is None or replica_group is None:
+                raise RuntimeError(
+                    "HYBRID_SHARD requires initialized target-TP and target-DP groups"
+                )
+            if dist.get_world_size(shard_group) <= 1:
+                raise ValueError("HYBRID_SHARD shard group must contain multiple ranks")
+            if dist.get_world_size(replica_group) <= 1:
+                raise ValueError(
+                    "HYBRID_SHARD replica group must contain multiple target islands"
+                )
+            fsdp_process_group = (shard_group, replica_group)
+            # Each replica owns an identical copy of one parameter shard after
+            # HSDP synchronization. Count that shard once, across the shard
+            # group, when computing the global norm.
+            grad_norm_process_group = shard_group
         return cls(
             world_size=dist.get_world_size(),
             tp_size=tp_size,
             sp_ulysses_size=sp_ulysses_size,
             sp_ring_size=sp_ring_size,
             sharding_strategy=sharding_strategy,
             param_dtype=param_dtype,
-            fsdp_process_group=dist.group.WORLD,
+            fsdp_process_group=fsdp_process_group,
+            reduction_process_group=dist.group.WORLD,
+            grad_norm_process_group=grad_norm_process_group,
             **handles,
         )
 
@@ -299,8 +328,11 @@ def set_optimizer(self, optimizer) -> None:
     def _configure_optimizer_grad_norm(self) -> None:
         configure = getattr(self.optimizer, "configure_grad_norm_reduction", None)
         if configure is not None:
+            process_group = self.parallel_config.grad_norm_process_group
+            if process_group is None:
+                process_group = self.parallel_config.fsdp_process_group
             configure(
-                process_group=self.parallel_config.fsdp_process_group,
+                process_group=process_group,
                 enabled=(
                     self._wrapped
                     and self.parallel_config.sharding_strategy != "NO_SHARD"

specforge/training/controller.py

@@ -439,12 +439,24 @@ def _accumulate_ratio_metrics(self, values: Dict[str, Any]) -> None:
                 denominator = previous[1] + denominator
             self._ratio_totals[name] = (numerator, denominator)
 
+    def _reduction_group(self):
+        """Group for loss/metric reductions across all sample-parallel ranks.
+
+        HSDP's fsdp_process_group is a (shard, replica) pair, so reductions use
+        the explicit WORLD-wide reduction group; every other topology falls
+        back to the FSDP group, which is a plain process group there.
+        """
+        parallel_config = getattr(self.backend, "parallel_config", None)
+        process_group = getattr(parallel_config, "reduction_process_group", None)
+        if process_group is None:
+            process_group = getattr(parallel_config, "fsdp_process_group", None)
+        return process_group
+
     def _normalize_gradients(self, local_denominator: torch.Tensor) -> None:
         import torch.distributed as dist
 
         denominator = local_denominator.clone()
-        parallel_config = getattr(self.backend, "parallel_config", None)
-        process_group = getattr(parallel_config, "fsdp_process_group", None)
+        process_group = self._reduction_group()
         world_size = 1
         if dist.is_available() and dist.is_initialized():
             world_size = dist.get_world_size(group=process_group)
@@ -477,8 +489,7 @@ def _result(
             if isinstance(out.loss, torch.Tensor)
             else torch.device("cpu")
         )
-        parallel_config = getattr(self.backend, "parallel_config", None)
-        process_group = getattr(parallel_config, "fsdp_process_group", None)
+        process_group = self._reduction_group()
         structured = _reduce_eagle3_metrics(
             out.metrics,
             device=metric_device,

tests/test_config/test_schema.py

@@ -135,6 +135,28 @@ def test_input_modality_is_extensible_and_resolved_by_providers(self):
             with self.subTest(modality=modality), self.assertRaises(ValidationError):
                 Config.model_validate(invalid)
 
+    def test_hybrid_shard_requires_multiple_colocated_target_islands(self):
+        payload = _online_payload("dspark")
+        payload["deployment"] = {
+            "mode": "local_colocated",
+            "trainer": {
+                "nnodes": 2,
+                "nproc_per_node": 8,
+                "master_addr": "trainer-0",
+            },
+        }
+        payload["training"].update({"tp_size": 8, "fsdp_sharding": "HYBRID_SHARD"})
+
+        config = Config.model_validate(payload)
+
+        self.assertEqual(config.training.fsdp_sharding, "HYBRID_SHARD")
+        config.validate_world_size(16)
+
+        payload["deployment"]["trainer"]["nnodes"] = 1
+        payload["deployment"]["trainer"].pop("master_addr")
+        with self.assertRaisesRegex(ValidationError, "two colocated target islands"):
+            Config.model_validate(payload)
+
     def test_target_output_sharding_is_unavailable_with_server_capture(self):
         payload = _online_payload()
         payload["model"]["shard_target_output"] = True

tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py

@@ -17,6 +17,28 @@ def _make_optimizer(seed=0, **kwargs):
 
 
 class TestClipGradNormSingleProcess(unittest.TestCase):
+    def test_backend_uses_the_hsdp_shard_group_for_grad_norm(self):
+        class RecordingOptimizer:
+            def configure_grad_norm_reduction(self, **kwargs):
+                self.config = kwargs
+
+        fsdp_groups = (object(), object())
+        shard_group = object()
+        backend = FSDPTrainingBackend(
+            ParallelConfig(
+                sharding_strategy="HYBRID_SHARD",
+                fsdp_process_group=fsdp_groups,
+                grad_norm_process_group=shard_group,
+            )
+        )
+        backend._wrapped = True
+        optimizer = RecordingOptimizer()
+
+        backend.set_optimizer(optimizer)
+
+        self.assertIs(optimizer.config["process_group"], shard_group)
+        self.assertTrue(optimizer.config["enabled"])
+
     def test_matches_torch_reference(self):
         model, optimizer = _make_optimizer()
         grad = torch.randn(8, 8)

tests/test_runtime/test_seam_fixes.py

@@ -91,6 +91,53 @@ def test_from_distributed_no_dist(self):
         self.assertEqual(pc.tp_size, 2)
         self.assertEqual(pc.sp_size, 2)
 
+    def test_hybrid_shard_separates_fsdp_metrics_and_grad_norm_groups(self):
+        tp_group = object()
+        dp_group = object()
+
+        def group_size(group=None):
+            if group is tp_group:
+                return 8
+            if group is dp_group:
+                return 2
+            return 16
+
+        with (
+            mock.patch(
+                "specforge.training.backend.dist.is_initialized", return_value=True
+            ),
+            mock.patch(
+                "specforge.training.backend.dist.get_world_size",
+                side_effect=group_size,
+            ),
+            mock.patch("specforge.distributed.get_tp_group", return_value=tp_group),
+            mock.patch("specforge.distributed.get_dp_group", return_value=dp_group),
+            mock.patch(
+                "specforge.distributed.get_draft_dp_group", return_value=object()
+            ),
+            mock.patch(
+                "specforge.distributed.get_draft_sp_group", return_value=object()
+            ),
+            mock.patch(
+                "specforge.distributed.get_sp_ulysses_group", return_value=object()
+            ),
+            mock.patch(
+                "specforge.distributed.get_sp_ring_group", return_value=object()
+            ),
+            mock.patch("specforge.distributed.get_device_mesh", return_value=object()),
+            mock.patch(
+                "specforge.distributed.get_tp_device_mesh", return_value=object()
+            ),
+        ):
+            pc = ParallelConfig.from_distributed(
+                tp_size=8,
+                sharding_strategy="HYBRID_SHARD",
+            )
+
+        self.assertEqual(pc.fsdp_process_group, (tp_group, dp_group))
+        self.assertIs(pc.reduction_process_group, torch.distributed.group.WORLD)
+        self.assertIs(pc.grad_norm_process_group, tp_group)
+
     def test_frozen_target_tables_are_ignored_by_fsdp(self):
         class Composite(nn.Module):
             def __init__(self):