fix: auto-compute dp_replicate_size from world_size (#1302)

## Summary
- When `dp_shard_size < world_size` (e.g., `dp_shard_size=4` on 8 GPUs
across 2 nodes), `ParallelismConfig` raises `total_size (4) does not
match num_processes (8)` because `dp_replicate_size` defaults to 1
- Auto-compute `dp_replicate_size = world_size // (dp_shard_size *
cp_size)` so intra-node FSDP2 sharding + inter-node data-parallel
replication works without manual config
- This enables `dp_shard_size` to be set to per-node GPU count (better
NVLink utilization) while automatically creating replicas across nodes

## Test plan
- [ ] Verify single-node training (dp_shard_size == world_size,
dp_replicate_size == 1) unchanged
- [ ] Verify multi-node with dp_shard_size < world_size creates correct
replica groups
- [ ] Verify existing EAGLE3/DFlash configs still work

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Refactor**
* Enhanced parallelism configuration initialization in the speculative
decoding example to better handle distributed training scenarios.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Ye Yu <yeyu@nvidia.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: yeyu-nvidia

examples/speculative_decoding/main.py

@@ -212,8 +212,23 @@ def train():
             "Either data.data_path or data.offline_data_path must be set in the config."
         )
     if training_args.cp_size > 1 or training_args.dp_shard_size > 1:
+        # Auto-compute dp_replicate_size so that
+        # dp_replicate_size * dp_shard_size * cp_size == world_size.
+        # Note: torch.cuda.device_count() returns per-node GPU count, not world_size.
+        # WORLD_SIZE (set by torchrun/accelerate) gives the correct multi-node total.
+        world_size = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count()))
+        parallel_size = training_args.dp_shard_size * training_args.cp_size
+        if world_size % parallel_size != 0:
+            raise ValueError(
+                f"world_size ({world_size}) must be divisible by "
+                f"dp_shard_size ({training_args.dp_shard_size}) * cp_size ({training_args.cp_size}) "
+                f"= {parallel_size}"
+            )
+        dp_replicate_size = world_size // parallel_size
         training_args.parallelism_config = ParallelismConfig(
-            cp_size=training_args.cp_size, dp_shard_size=training_args.dp_shard_size
+            cp_size=training_args.cp_size,
+            dp_shard_size=training_args.dp_shard_size,
+            dp_replicate_size=dp_replicate_size,
         )
     if training_args.cp_size > 1:
         patch_ring_attention_for_ttt()