fix: add divisibility guard and clarify WORLD_SIZE fallback

Address review feedback:
- Add ValueError if world_size is not divisible by dp_shard_size * cp_size
- Comment that torch.cuda.device_count() is per-node, not world_size

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

Author: yeyu-nvidia

examples/speculative_decoding/main.py

@@ -214,8 +214,16 @@ def train():
     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,