From 09e88e021392d1658c81f1395430824975c04ca0 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 6 Jul 2026 22:30:42 +0000 Subject: [PATCH] fsdp: assert weight-sync grouping covers the train world connect_rollout_engines' per-engine rank span silently disagreed with rollout.init_rollout_engines' (raw rollout_num_gpus_per_engine vs min(rollout_num_gpus_per_engine, num_gpus_per_node)). For every current single-node config the two coincide and behavior is unchanged. For multi-node colocate (engine TP spanning nodes) NEITHER convention yields a working weight push: the FromTensor path is CUDA-IPC based and cannot cross nodes, so the old code failed much later inside update_bucket_weights. The new assert (world == engines x span) makes such configs fail at connect time with the three numbers spelled out. Proper multi-node weight sync needs per-node-shard grouping and is out of scope here. Co-Authored-By: Claude Fable 5 --- .../fsdp_utils/diffusion_update_weight_utils.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index dc6193c8..c98114ea 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -107,10 +107,18 @@ def connect_rollout_engines( ) -> None: self.rollout_engines = rollout_engines + # Match rollout.init_rollout_engines' per-engine span; a mismatch would + # leave orphaned ranks that fail much later in update_bucket_weights. + num_gpu_per_engine = min(self.args.rollout_num_gpus_per_engine, self.args.num_gpus_per_node) + assert dist.get_world_size() == len(rollout_engines) * num_gpu_per_engine, ( + f"train world {dist.get_world_size()} != {len(rollout_engines)} engines x " + f"{num_gpu_per_engine} gpus per engine" + ) + # Here we assume the gpu id of rollout engines and train actors are the same. for i, engine in enumerate(self.rollout_engines): - start_rank = i * self.args.rollout_num_gpus_per_engine - end_rank = (i + 1) * self.args.rollout_num_gpus_per_engine + start_rank = i * num_gpu_per_engine + end_rank = (i + 1) * num_gpu_per_engine group_ranks = list(range(start_rank, end_rank)) new_group = dist.new_group( ranks=group_ranks,