From ebe0558b5e19f3496ef0cb70ed6c712a934183e2 Mon Sep 17 00:00:00 2001 From: wdykas Date: Mon, 22 Jun 2026 06:11:50 -0700 Subject: [PATCH 1/2] Refit into multiple destination pools; tied-embedding + UVM fixes Generic refit/UVM enablers, independent of any disaggregation code: - refit.py: swap_model_weights gains num_dst_pools/dst_pool_index to refit the source model into N disjoint destination pools (one collective pass per pool). The reshard-plan cache is keyed by pool index so a source-only rank does not alias plans across pools, and each pool's plan is built once then reused across refits in lockstep. Default (1, 0) reproduces single-pool behavior exactly. - planner.py: when the destination materializes output_layer.weight but the source has tied embeddings (e.g. pp=1 -> pp=2), source it from embedding.word_embeddings.weight (identical shape + vocab/TP shard). - unified_memory.py: reject the UVM mempool up front when PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (which torch.cuda.MemPool cannot coexist with) so callers fall back to plain allocation instead of crashing mid-buffer-setup. Single-pool refit is covered by the existing test_model_swap; the multi-pool path is exercised end-to-end by the disaggregated RL refit in a later change. Co-Authored-By: Claude Opus 4.8 Signed-off-by: wdykas --- megatron/core/inference/unified_memory.py | 10 ++++ megatron/core/resharding/planner.py | 14 +++++ megatron/core/resharding/refit.py | 64 ++++++++++++++++------- 3 files changed, 69 insertions(+), 19 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 53d8d862b92..613fb246f10 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -277,6 +277,16 @@ def create_unified_mempool() -> "MemPool": + details ) else: + # torch.cuda.MemPool can't coexist with expandable_segments; bail so the + # caller falls back to non-UVM allocation. + alloc_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "") + if "expandable_segments:true" in alloc_conf.lower(): + raise UnifiedMemoryUnsupportedError( + "UVM mempool is incompatible with the expandable-segments allocator " + "(PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, which " + "torch.cuda.MemPool does not support). Unset expandable_segments to " + "use UVM." + ) return MemPool(allocator=_alloc) diff --git a/megatron/core/resharding/planner.py b/megatron/core/resharding/planner.py index f0e38004d0e..dabf822a298 100644 --- a/megatron/core/resharding/planner.py +++ b/megatron/core/resharding/planner.py @@ -404,6 +404,20 @@ def _extract_metadata(module, rank_offset): dst_rank_params = dst_param_metadata_by_rank.get(dst_rank, {}) for resolved_name, dst_metadata in dst_rank_params.items(): src_meta_list = src_param_metadata.get(resolved_name) + if not src_meta_list and resolved_name.endswith("output_layer.weight"): + # Tied embeddings: the source shares the output projection with + # the input embedding, so it has no separate output_layer.weight. + # A pp>1 destination materializes one (embedding and output land + # on different stages), e.g. pp=1 (tied) -> pp=2. Source it from + # the embedding weight (same shape + vocab/TP shard); that tensor + # then feeds both the destination embedding and output_layer. + for emb_name in ( + "embedding.word_embeddings.weight", + "word_embeddings.weight", + ): + src_meta_list = src_param_metadata.get(emb_name) + if src_meta_list: + break if not src_meta_list: raise RuntimeError( f"Destination parameter '{resolved_name}' on rank {dst_rank} " diff --git a/megatron/core/resharding/refit.py b/megatron/core/resharding/refit.py index 36ba914d33b..624ffec58d0 100644 --- a/megatron/core/resharding/refit.py +++ b/megatron/core/resharding/refit.py @@ -49,6 +49,9 @@ class _PlanCacheKey: # global ranks. src_rank_offset: int = 0 dst_rank_offset: int = 0 + # Multi-pool refit: keep each destination pool's plan distinct so source-only + # ranks (dst_config=None on every pool) don't alias them. See swap_model_weights. + pool_index: int = 0 def _get_config_tuple(core) -> Optional[Tuple[int, int, int, int, int]]: @@ -83,6 +86,7 @@ def _build_plan_cache_key( group=None, src_rank_offset: int = 0, dst_rank_offset: int = 0, + pool_index: int = 0, ) -> _PlanCacheKey: """Build cache key for reshard plan.""" # group.rank() supports cross-cluster ProcessGroups. @@ -94,6 +98,7 @@ def _build_plan_cache_key( num_experts=num_experts, src_rank_offset=src_rank_offset, dst_rank_offset=dst_rank_offset, + pool_index=pool_index, ) @@ -191,7 +196,9 @@ def _unwrap_model_cores(src_model, target_model): return src_core, tgt_core, num_experts -def _build_or_get_plan(src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset): +def _build_or_get_plan( + src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset, pool_index=0 +): """Return the cached reshard plan, building it (collectively) if not yet cached. All participating ranks must call this simultaneously when the plan is not @@ -205,6 +212,7 @@ def _build_or_get_plan(src_core, tgt_core, num_experts, group, src_rank_offset, group=group, src_rank_offset=src_rank_offset, dst_rank_offset=dst_rank_offset, + pool_index=pool_index, ) if cache_key not in _plan_cache: _plan_cache[cache_key] = build_centralized_reshard_plan( @@ -326,6 +334,8 @@ def swap_model_weights( src_rank_offset: int = 0, dst_rank_offset: int = 0, transform: Optional[ReshardTransform] = None, + num_dst_pools: int = 1, + dst_pool_index: int = 0, ): """ Orchestrate weight swap/refit. @@ -345,6 +355,12 @@ def swap_model_weights( transform: Optional ReshardTransform for custom format conversion. If None, the cached transform (from prepare_swap_model_weights) is used automatically when the receiver needs MXFP8 conversion. + num_dst_pools / dst_pool_index: refit into ``num_dst_pools`` disjoint + destination pools (e.g. disaggregated prefill/decode instances on + separate rank windows), one collective pass per pool. This rank + receives into ``target_model`` only on its own pool's pass + (``pool == dst_pool_index``) and is a pure source otherwise. + Defaults ``(1, 0)`` reproduce the single-destination behavior. """ if isinstance(refit_method, str): service = get_or_create_service(refit_method, group=group) @@ -353,24 +369,33 @@ def swap_model_weights( else: raise TypeError("refit_method must be a str backend name or a CopyService instance") - # Auto-resolve MXFP8 transform from the cached plan when no - # explicit transform was provided. - if transform is None: - src_core, tgt_core, num_experts = _unwrap_model_cores(src_model, target_model) - plan = _build_or_get_plan( - src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset + for pool in range(num_dst_pools): + target = target_model if pool == dst_pool_index else None + + # The plan-build is collective. Pass ``pool`` so a source-only rank + # (target=None, same cache key every pool) doesn't cache-hit and skip the + # collective on a later pass while target ranks run it -> deadlock. Each + # pool's plan is then built once and reused across refits in lockstep. + # Auto-resolve MXFP8 transform from the cached plan when no explicit + # transform was provided (re-resolved per pool: the target differs). + pass_transform = transform + if pass_transform is None: + src_core, tgt_core, num_experts = _unwrap_model_cores(src_model, target) + plan = _build_or_get_plan( + src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset, pool + ) + pass_transform = plan.transform + + reshard_model_weights( + src_model, + target, + service=service, + group=group, + src_rank_offset=src_rank_offset, + dst_rank_offset=dst_rank_offset, + transform=pass_transform, + pool_index=pool, ) - transform = plan.transform - - reshard_model_weights( - src_model, - target_model, - service=service, - group=group, - src_rank_offset=src_rank_offset, - dst_rank_offset=dst_rank_offset, - transform=transform, - ) def _harmonize_buffer_dtypes(plan, src_core, tgt_core, group=None): @@ -431,6 +456,7 @@ def reshard_model_weights( src_rank_offset: int = 0, dst_rank_offset: int = 0, transform: Optional[ReshardTransform] = None, + pool_index: int = 0, ): """Reshard and copy model weights from ``src_model`` to ``target_model`` using ``service``. @@ -448,7 +474,7 @@ def reshard_model_weights( """ src_core, tgt_core, num_experts = _unwrap_model_cores(src_model, target_model) plan = _build_or_get_plan( - src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset + src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset, pool_index ) _harmonize_buffer_dtypes(plan, src_core, tgt_core, group=group) execute_reshard_plan( From a4b29242eeef64c7f09dc200a8430af9064df9f5 Mon Sep 17 00:00:00 2001 From: wdykas Date: Mon, 22 Jun 2026 06:11:50 -0700 Subject: [PATCH 2/2] lint Signed-off-by: wdykas --- megatron/core/resharding/planner.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/megatron/core/resharding/planner.py b/megatron/core/resharding/planner.py index dabf822a298..1eda91e914b 100644 --- a/megatron/core/resharding/planner.py +++ b/megatron/core/resharding/planner.py @@ -411,10 +411,7 @@ def _extract_metadata(module, rank_offset): # on different stages), e.g. pp=1 (tied) -> pp=2. Source it from # the embedding weight (same shape + vocab/TP shard); that tensor # then feeds both the destination embedding and output_layer. - for emb_name in ( - "embedding.word_embeddings.weight", - "word_embeddings.weight", - ): + for emb_name in ("embedding.word_embeddings.weight", "word_embeddings.weight"): src_meta_list = src_param_metadata.get(emb_name) if src_meta_list: break