diff --git a/awex/converter/mcore_converter.py b/awex/converter/mcore_converter.py index d4c6da7..ccdfec9 100644 --- a/awex/converter/mcore_converter.py +++ b/awex/converter/mcore_converter.py @@ -651,7 +651,10 @@ def _convert_expert_bias_param( ) -> Tuple[str, torch.Tensor]: """Convert bias parameters""" if "expert_bias" in name: - return ("mlp.gate.expert_bias", parameter.to(torch.bfloat16)) + # SGLang keeps the auxiliary router expert bias in fp32 even when + # router matmul runs in bf16. Downcasting here makes AWEX-loaded + # SGLang diverge from the native HF load at step 0. + return ("mlp.gate.expert_bias", parameter.to(torch.float32)) else: raise NotImplementedError(f"Unsupported bias parameter name: {name}") @@ -745,9 +748,22 @@ def __init__( tf_config: TransformerConfig, ): super().__init__(hf_config, rank_info, infer_conf, tf_config=tf_config) - self.layer_group_size = int(_cfg_get(tf_config, "layer_group_size", 1) or 1) - self.linear_attn_norm_group_size = _cfg_get( - tf_config, "linear_attn_norm_group_size", None + # `layer_group_size` / `linear_attn_norm_group_size` are BailingMoeV2.5 + # hybrid-attention params. AReaL's MegatronEngine does not inject them into + # the Megatron TransformerConfig, so fall back to the HF config (matching + # AReaL's own bailing_moe.py mapping: linear_attn_norm_group_size defaults + # to HF `group_norm_size`). Without this, layer_group_size stays 1 and every + # Lightning-attention layer is misrouted to the MLA path -> `linear_gate` + # raises "Unsupported parameter name". + self.layer_group_size = int( + _cfg_get(tf_config, "layer_group_size", None) + or getattr(hf_config, "layer_group_size", None) + or 1 + ) + self.linear_attn_norm_group_size = ( + _cfg_get(tf_config, "linear_attn_norm_group_size", None) + or getattr(hf_config, "linear_attn_norm_group_size", None) + or getattr(hf_config, "group_norm_size", None) ) self.fuse_qkv_a_proj = getattr(hf_config, "q_lora_rank", None) is not None self.qkv_a_proj_cache: Dict[str, Dict[str, torch.Tensor]] = {} @@ -763,11 +779,15 @@ def _is_linear_layer(self, layer_number: int) -> bool: return (layer_number + 1) % self.layer_group_size != 0 def _convert_g_norm_weight(self, parameter: torch.Tensor) -> torch.Tensor: - if not self.linear_attn_norm_group_size: - return parameter.clone().detach() - group_size = int(self.linear_attn_norm_group_size) - tp_size = max(int(self.rank_info.tp_size), 1) - return parameter.clone().detach().reshape(group_size // tp_size, -1) + # SGLang's BailingGroupRMSNormGate stores g_norm.weight as a flat 1D + # tensor [hidden_inner_size // tp_size] and applies grouping at runtime. + # The Megatron weight is also 1D; reshaping it to 2D here only added a + # leading dim (data order unchanged) and made the train-side metadata + # ndim (2) differ from the infer-side ndim (1), which crashed the AWEX + # transfer plan with `IndexError: tuple index out of range` in + # _build_region_communication_plan (Problem 38). Keep it 1D to match the + # SGLang canonical layout. + return parameter.clone().detach().reshape(-1) def _convert_lightning_attention_param( self, name: str, parameter: torch.Tensor, layer_number: str @@ -776,6 +796,10 @@ def _convert_lightning_attention_param( return [] if "self_attention.pre_gate_norm.weight" in name: return [("attention.g_norm.weight", self._convert_g_norm_weight(parameter))] + # BailingMoeV2.5 uses "gate_norm" instead of "pre_gate_norm" for + # the Lightning Attention gating normalization layer. + if "self_attention.gate_norm.weight" in name: + return [("attention.g_norm.weight", self._convert_g_norm_weight(parameter))] name_mapping = { "self_attention.input_layernorm.weight": "input_layernorm.weight", @@ -817,6 +841,17 @@ def _try_fuse_qkv_a_proj( if other_key not in layer_cache: return [] + # LinearMLAShardingMixin.get_sharding_strategy() always declares + # ``attention.fused_qkv_a_proj_with_mqa.weight`` as NO_SHARDING, so the + # transfer plan ships the converted tensor UNCUT to every SGLang TP + # rank. The receiver's load_weights layout is the plain concatenation + # ``torch.cat([q_a, kv_a], dim=0)``; any infer-TP interleaved layout + # ([q_0;kv_0;q_1;kv_1;...]) scrambles + # rows for every recipient. Both inputs were already gathered to full + # tensors via get_full_tensor above, so the canonical concat is correct + # regardless of the train attn-TP size. Keep it unconditional unless + # the sharding metadata and transfer plan are changed to perform a + # real TP split of this parameter. fused_tensor = torch.cat( [layer_cache["q_a_proj"], layer_cache["kv_a_proj"]], dim=0 ) @@ -855,9 +890,27 @@ def _convert_mla_attention_param( return super()._convert_attention_param(name, parameter, layer_number) + # MLA-specific parameter name fragments that uniquely identify MLA layers, + # even when the global layer ID is unavailable (PP meta-resolution phase). + _MLA_PARAM_MARKERS = ( + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_down_proj", + "linear_kv_up_proj", + "linear_q_proj", + ) + + def _is_mla_param(self, name: str) -> bool: + return any(marker in name for marker in self._MLA_PARAM_MARKERS) + def _convert_attention_param( self, name: str, parameter: torch.Tensor, layer_number: str ) -> List[Tuple[str, torch.Tensor]]: + # When PP stage layer ID map is unavailable (during meta resolution), + # local layer IDs may not reflect the true global position. Fall back + # to detecting layer type from the parameter name itself. + if self._is_mla_param(name): + return self._convert_mla_attention_param(name, parameter, layer_number) if self._is_linear_layer(int(layer_number)): return self._convert_lightning_attention_param( name, parameter, layer_number @@ -1001,7 +1054,23 @@ def transform_mcore_qkv_weight(weight: torch.Tensor, tf_config: TransformerConfi hidden_size * 2 + 2 * hidden_size ) # Q extends to 2, k v use hidden size - if actual_size == expected_replicated_size: + if actual_size == expected_compact_size: + # Compact GQA/MHA interleaved format (Megatron default): + # [Q0,K0,V0, Q1,K1,V1, ...] chunked by num_kv_heads. + # Must check BEFORE replicated because MHA (num_heads == num_kv_heads) + # makes compact_size == replicated_size. + query_list = [] + key_list = [] + value_list = [] + for qkv in torch.chunk(weight, total_num_kv_heads, dim=0): + q, k, v = qkv.split([each_query_size, each_kv_size, each_kv_size], dim=0) + query_list.append(q) + key_list.append(k) + value_list.append(v) + all_query = torch.cat(query_list, dim=0) + all_key = torch.cat(key_list, dim=0) + all_value = torch.cat(value_list, dim=0) + elif actual_size == expected_replicated_size: # Replicated format: K and V are replicated to match query heads # Split into Q, K, V where each has size hidden_size q, k, v = weight.split([hidden_size, hidden_size, hidden_size], dim=0) @@ -1023,21 +1092,6 @@ def transform_mcore_qkv_weight(weight: torch.Tensor, tf_config: TransformerConfi all_key = torch.cat(k_unique, dim=0).reshape(-1, k.shape[-1]) all_value = torch.cat(v_unique, dim=0).reshape(-1, v.shape[-1]) all_query = q - - elif actual_size == expected_compact_size: - # Compact GQA format: K and V have reduced size - query_list = [] - key_list = [] - value_list = [] - for qkv in torch.chunk(weight, total_num_kv_heads, dim=0): - q, k, v = qkv.split([each_query_size, each_kv_size, each_kv_size], dim=0) - query_list.append(q) - key_list.append(k) - value_list.append(v) - # concat the query, key, value - all_query = torch.cat(query_list, dim=0) - all_key = torch.cat(key_list, dim=0) - all_value = torch.cat(value_list, dim=0) elif actual_size == expected_qwen3_gqa_size: query_list = [] key_list = [] @@ -1143,7 +1197,26 @@ def transform_mcore_qkv_bias(bias: torch.Tensor, tf_config: TransformerConfig): expected_compact_size = (each_query_size + 2 * each_kv_size) * total_num_kv_heads expected_replicated_size = 3 * hidden_size # Q, K, V all have full hidden_size - if actual_size == expected_replicated_size: + if actual_size == expected_compact_size: + # Compact GQA/MHA interleaved format (Megatron default): + # [Q0,K0,V0, Q1,K1,V1, ...] chunked by num_kv_heads. + # Must check BEFORE replicated because MHA (num_heads == num_kv_heads) + # makes compact_size == replicated_size, and a fused QKV linear uses + # the same output-row ordering for weight and bias — this branch order + # mirrors transform_mcore_qkv_weight so both stay consistent. + query_list = [] + key_list = [] + value_list = [] + for qkv in torch.chunk(bias, total_num_kv_heads, dim=0): + q, k, v = qkv.split([each_query_size, each_kv_size, each_kv_size], dim=0) + query_list.append(q) + key_list.append(k) + value_list.append(v) + # concat the query, key, value + all_query = torch.cat(query_list, dim=0) + all_key = torch.cat(key_list, dim=0) + all_value = torch.cat(value_list, dim=0) + elif actual_size == expected_replicated_size: # Replicated format: K and V are replicated to match query heads # Split into Q, K, V where each has size hidden_size q, k, v = bias.split([hidden_size, hidden_size, hidden_size], dim=0) @@ -1165,21 +1238,6 @@ def transform_mcore_qkv_bias(bias: torch.Tensor, tf_config: TransformerConfig): all_key = torch.cat(k_unique, dim=0).reshape(-1) all_value = torch.cat(v_unique, dim=0).reshape(-1) all_query = q - - elif actual_size == expected_compact_size: - # Compact GQA format: K and V have reduced size - query_list = [] - key_list = [] - value_list = [] - for qkv in torch.chunk(bias, total_num_kv_heads, dim=0): - q, k, v = qkv.split([each_query_size, each_kv_size, each_kv_size], dim=0) - query_list.append(q) - key_list.append(k) - value_list.append(v) - # concat the query, key, value - all_query = torch.cat(query_list, dim=0) - all_key = torch.cat(key_list, dim=0) - all_value = torch.cat(value_list, dim=0) else: raise ValueError( f"QKV bias size mismatch - unsupported format:\n" diff --git a/awex/models/ling.py b/awex/models/ling.py index 2c7549e..c147089 100644 --- a/awex/models/ling.py +++ b/awex/models/ling.py @@ -37,6 +37,18 @@ class BailingMoeShardingStrategy(ShardingStrategy): def get_sharding_strategy(self, parameter_name, **kwargs): if self.engine_name == "mcore": if "query_key_value" in parameter_name: + # convert_qkv_weight_along_tp_attention all-gathers qkv across + # the train attn-TP group, repacks it into the SGLang fused + # layout, then hands each rank only shards[attn_tp_rank] when + # attn_tp_size > 1 — the converted tensor is TP-sharded along + # dim 0, not replicated. Declaring NO_SHARDING here made every + # train rank a full replica whose declared extent covered only + # rows [0, N/tp), so the transfer plan generated ops for infer + # tp_rank 0 only and Lightning qkv on tp_rank > 0 stayed + # all-zero. + attn_tp_size = self.rank_info.attn_tp_size + if attn_tp_size > 1: + return ShardingType.TP_SHARDING, 0, attn_tp_size return ShardingType.NO_SHARDING, 0, 1 return super().get_sharding_strategy(parameter_name, **kwargs) diff --git a/awex/models/ling_linear.py b/awex/models/ling_linear.py index f2fea9e..b2502ac 100644 --- a/awex/models/ling_linear.py +++ b/awex/models/ling_linear.py @@ -99,7 +99,17 @@ class SGlangToHFWeightConverterBailingMoeLinear( LinearMLASGlangConverterMixin, SGlangToHFWeightConverter, ): - pass + def convert_param( + self, name: str, parameter: torch.Tensor + ) -> List[Tuple[str, torch.Tensor]]: + # SGLang's BailingMoe names the input embedding ``word_embeddings``, + # while the canonical HF name (and the train-side mcore converter output) + # is ``embed_tokens``. Normalize here so the transfer-plan key sets on + # both sides align; the base converter has no rule for ``word_embeddings`` + # and would otherwise pass it through unchanged. + if name == "model.word_embeddings.weight": + name = "model.embed_tokens.weight" + return super().convert_param(name, parameter) CONFIG = [ diff --git a/awex/reader/nccl_reader.py b/awex/reader/nccl_reader.py index e576b60..c548bd0 100644 --- a/awex/reader/nccl_reader.py +++ b/awex/reader/nccl_reader.py @@ -463,7 +463,17 @@ def _update_weights_in_colocate_mode(self, step_id, **kwargs): print_current_gpu_status( f"after weights update using NCCL for rank {self.rank_coordinate}" ) + # Dropping the reference is NOT enough: cudaIpcCloseMemHandle runs in + # the tensor deleter, which only fires once refcounts/GC actually + # release the IPC-imported tensors. If that close lags past the train + # side's release+empty_cache+realloc (train acts right after + # weights_update_finished), the stale IPC mapping overlaps train's + # fresh allocations and either side faults with an illegal memory + # access at a drifting location (in reload, resume, or first read). + # Force the close to complete BEFORE signalling the train side. self.deserialized_weights = None + gc.collect() + device_util.synchronize() duration = time.time() - start_time compute_statistics( self._history_update_weights_time, diff --git a/awex/sharding/param_sharding.py b/awex/sharding/param_sharding.py index 0928164..fb8b836 100644 --- a/awex/sharding/param_sharding.py +++ b/awex/sharding/param_sharding.py @@ -267,7 +267,11 @@ def get_sharding_strategy(self, parameter_name, **kwargs): return ShardingType.NO_SHARDING, 0, 1 if "norm" in parameter_name: return ShardingType.NO_SHARDING, 0, 1 - if "embedding" in parameter_name: + if "embedding" in parameter_name or "embed_tokens" in parameter_name: + # ``embed_tokens`` is the canonical HF name for the input embedding + # (mcore/sglang converters both normalize to it); it must inherit + # embedding sharding semantics even though it does not contain the + # literal substring "embedding". return self.get_embedding_sharding_strategy(parameter_name, **kwargs) if "lm_head" in parameter_name: return self.get_lm_head_sharding_strategy(parameter_name, **kwargs) diff --git a/awex/tests/test_attn_tp_sharding_cp_bug.py b/awex/tests/test_attn_tp_sharding_cp_bug.py new file mode 100644 index 0000000..a83eea6 --- /dev/null +++ b/awex/tests/test_attn_tp_sharding_cp_bug.py @@ -0,0 +1,132 @@ +# Copyright (c) Ant Group. Licensed under the Apache License, Version 2.0. +""" +Regression test for attention qkv sharding declarations under +**train attn_tp == 1 (pure CP) + infer TP > 1** topologies. + +Background: + With a colocate allocation where the train side uses pure context + parallelism for attention (e.g. attn=d2p4c8, attn_tp_size == 1), + inference collapsed at the first post-update step (~98% rejected + generations), while an otherwise identical run with attn=d1p4t4c2 + (attn_tp_size == 4) was healthy. The only differing variable was + whether the train attention used TP. + + The sharding-strategy comment for query_key_value in ling.py records + the same failure mode: + > Declaring NO_SHARDING ... the transfer plan generated ops for infer + > tp_rank 0 only and Lightning qkv on tp_rank > 0 stayed all-zero. + That fix only covers the `attn_tp_size > 1` branch (returns + TP_SHARDING, aligned with infer TP). When train `attn_tp_size == 1` + the strategy still falls back to NO_SHARDING -- but the inference side + runs TP > 1, so the transfer plan again only covers infer tp_rank 0 + and the qkv weights on tp_rank > 0 stay all-zero. + +These tests pin down the trigger condition as a regression case: + - attn_tp == 4 (healthy topology) -> query_key_value declared + TP_SHARDING, aligned with infer TP + - attn_tp == 1 (broken topology) -> query_key_value falls back to + NO_SHARDING (the trigger condition above) + +Note: this test asserts the trigger condition at the sharding-strategy +layer (no TP -> NO_SHARDING declaration). A full end-to-end test that +the transfer plan misses ops for infer tp_rank > 0 requires the +meta-resolver + plan-generation layers (TODO); the runtime failure mode +was cross-validated with an lr=0 run where the post-update importance +weights diverged from 1. +""" + +import pytest + +from awex.sharding.param_sharding import ShardingType +from awex.sharding.rank_info import RankInfo + + +def _make_rank_info(attn_tp_size: int, cp_size: int) -> RankInfo: + """Build a train-side (mcore) RankInfo. + + attn_tp_size == 1 with cp_size > 1 reproduces the broken pure-CP + topology (e.g. d2p4c8). + """ + return RankInfo( + tp_rank=0, + tp_size=attn_tp_size, + pp_rank=0, + pp_size=1, + dp_size=1, + dp_rank=0, + ep_rank=0, + ep_size=1, + ep_tp_rank=0, + ep_tp_size=1, + attn_tp_rank=0, + attn_tp_size=attn_tp_size, + attn_dp_rank=0, + world_size=max(attn_tp_size, 1) * max(cp_size, 1), + global_rank=0, + local_rank=0, + engine_rank=0, + is_infer=False, + cp_rank=0, + cp_size=cp_size, + cp_mode="ring" if cp_size > 1 else "none", + ) + + +def _make_strategy(attn_tp_size: int, cp_size: int): + from awex.models.ling_linear import BailingLinearMoeShardingStrategy + + return BailingLinearMoeShardingStrategy( + engine_name="mcore", + enable_dp_attention=False, + enable_dp_lm_head=False, + moe_dense_tp_size=1, + tp_size=attn_tp_size, + ep_size=1, + ep_tp_size=1, + rank_info=_make_rank_info(attn_tp_size, cp_size), + ) + + +QKV_PARAM = "model.layers.0.attention.query_key_value.weight" + + +def test_qkv_sharding_with_attn_tp4_is_tp_sharded(): + """Healthy topology (attn_tp=4): TP_SHARDING, aligned with infer TP.""" + strat = _make_strategy(attn_tp_size=4, cp_size=2) + sharding_type, _dim, num_shards = strat.get_sharding_strategy(QKV_PARAM) + assert sharding_type == ShardingType.TP_SHARDING + assert num_shards == 4 + + +def test_qkv_sharding_with_attn_tp1_falls_back_to_no_sharding(): + """Broken topology (attn_tp=1, pure CP8): falls back to NO_SHARDING. + + This is the trigger condition: the train side declares a full + replica (NO_SHARDING) while the inference side runs TP > 1, so the + transfer plan only generates ops for infer tp_rank 0 and the qkv + weights on tp_rank > 0 stay all-zero. The TP_SHARDING fix only + covers the attn_tp > 1 branch, not this attn_tp == 1 + infer_tp > 1 + combination. + """ + strat = _make_strategy(attn_tp_size=1, cp_size=8) + sharding_type, _dim, num_shards = strat.get_sharding_strategy(QKV_PARAM) + # Current awex behavior (the trigger condition): NO_SHARDING fallback. + assert sharding_type == ShardingType.NO_SHARDING + assert num_shards == 1 + + +def test_cp_size_does_not_change_qkv_strategy(): + """Control: with attn_tp=1 fixed, cp_size 1 -> 8 does not change the decision. + + The attention sharding decision is entirely CP-agnostic: the CP + redundant replicas get no special handling at the strategy layer, + leaving the downstream collection/transfer layers responsible -- and + those layers under-transfer when attn_tp == 1 and infer_tp > 1. + """ + s_cp1 = _make_strategy(attn_tp_size=1, cp_size=1).get_sharding_strategy(QKV_PARAM) + s_cp8 = _make_strategy(attn_tp_size=1, cp_size=8).get_sharding_strategy(QKV_PARAM) + assert s_cp1 == s_cp8 # cp_size does not affect the attn sharding decision + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/awex/tests/test_nccl_stream_batch.py b/awex/tests/test_nccl_stream_batch.py new file mode 100644 index 0000000..1d66136 --- /dev/null +++ b/awex/tests/test_nccl_stream_batch.py @@ -0,0 +1,62 @@ +# Licensed to the Awex developers under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import torch + +from awex.transfer.nccl_stream_batch import ( + _clone_p2p_send_tensor, + _prepare_p2p_recv_tensor, + _sync_p2p_recv_tensor_pairs, +) + + +def test_prepare_p2p_recv_tensor_uses_dense_buffer_for_noncontiguous_view(): + base = torch.zeros(4, 4) + view = base[:, 1] + assert not view.is_contiguous() + + recv_tensor, copyback_pair = _prepare_p2p_recv_tensor(view) + + assert recv_tensor.is_contiguous() + assert copyback_pair is not None + recv_tensor.copy_(torch.arange(4, dtype=base.dtype)) + + _sync_p2p_recv_tensor_pairs([copyback_pair]) + + torch.testing.assert_close(base[:, 1], torch.arange(4, dtype=base.dtype)) + torch.testing.assert_close(base[:, 0], torch.zeros(4)) + + +def test_prepare_p2p_recv_tensor_reuses_contiguous_tensor(): + tensor = torch.zeros(4) + + recv_tensor, copyback_pair = _prepare_p2p_recv_tensor(tensor) + + assert recv_tensor is tensor + assert copyback_pair is None + + +def test_clone_p2p_send_tensor_returns_contiguous_clone(): + base = torch.arange(16, dtype=torch.float32).reshape(4, 4) + view = base[:, 1] + assert not view.is_contiguous() + + cloned = _clone_p2p_send_tensor(view) + + assert cloned.is_contiguous() + assert cloned.data_ptr() != view.data_ptr() + torch.testing.assert_close(cloned, view) diff --git a/awex/transfer/nccl_comm.py b/awex/transfer/nccl_comm.py index 7cac8f0..10b37b8 100644 --- a/awex/transfer/nccl_comm.py +++ b/awex/transfer/nccl_comm.py @@ -222,6 +222,7 @@ def update_weights_in_colocate_mode( ) +@torch.no_grad() def execute_tensors_to_copy(tensors_to_copy, copy_ops, recv_parameters, stage: str): start_time = time.time() num_ops = len(copy_ops) diff --git a/awex/transfer/nccl_stream_batch.py b/awex/transfer/nccl_stream_batch.py index a9b6ce7..82e4eeb 100644 --- a/awex/transfer/nccl_stream_batch.py +++ b/awex/transfer/nccl_stream_batch.py @@ -36,6 +36,36 @@ hang_detector = ThreadPoolExecutor(max_workers=1) +def _clone_p2p_send_tensor(tensor: torch.Tensor) -> torch.Tensor: + """Return a dense tensor suitable for torch.distributed P2P send.""" + return tensor.clone(memory_format=torch.contiguous_format) + + +def _prepare_p2p_recv_tensor( + tensor: torch.Tensor, +) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: + """Return a dense recv buffer and an optional copyback pair.""" + if tensor.is_contiguous(): + return tensor, None + recv_buffer = torch.empty( + tuple(tensor.shape), + dtype=tensor.dtype, + device=tensor.device, + ) + return recv_buffer, (tensor, recv_buffer) + + +@torch.no_grad() +def _sync_p2p_recv_tensor_pairs( + recv_tensor_pairs: list[tuple[torch.Tensor, torch.Tensor]], +) -> None: + if not recv_tensor_pairs: + return + for original_tensor, recv_buffer in recv_tensor_pairs: + original_tensor.copy_(recv_buffer) + recv_tensor_pairs.clear() + + class NcclColocateStreamBatchTransport: MAX_STREAMS = 64 @@ -82,12 +112,40 @@ def update_weights_in_colocate_mode( f"num_sends {num_sends}, num_recvs {num_recvs}" ) + chunk_mb = int(os.environ.get("AWEX_CHUNK_MB", "0") or "0") + + if chunk_mb > 0: + self._run_chunked( + task_id=task_id, + step_id=step_id, + train_to_infer_device_mapping=train_to_infer_device_mapping, + infer_to_train_device_mapping=infer_to_train_device_mapping, + transfer_rank=transfer_rank, + rank_coordinate=rank_coordinate, + world_size=world_size, + send_ops=send_ops, + recv_ops=recv_ops, + recv_transfer_plan=recv_transfer_plan, + weights_update_group=weights_update_group, + send_parameters=send_parameters, + recv_parameters=recv_parameters, + async_op=async_op, + chunk_bytes=chunk_mb * 1024 * 1024, + ) + duration = time.time() - start_time + logger.info( + f"Finished CHUNKED weights update for {task_id}, took {duration:.4f}s " + f"(chunk_mb={chunk_mb})" + ) + return + + # === LEGACY PATH (one-shot clone-then-transfer) === # Build P2P operations with sliced tensors all_send_p2p_ops = {} # peer_rank -> List[(plan_op, p2p_op)] all_recv_p2p_ops = {} # peer_rank -> List[(plan_op, p2p_op)] tensors_to_copy = [] + recv_tensor_pairs = [] train_slice_context = {} - non_contiguous_tensor_pairs = [] # Process send operations for peer_rank, ops in send_ops.items(): @@ -113,9 +171,18 @@ def update_weights_in_colocate_mode( recv_rank = train_to_infer_device_mapping.get( op.recv_rank, op.recv_rank ) + cloned = _clone_p2p_send_tensor(tensor_sliced) + # Wire-size parity with the receiver's dtype (see the + # chunked path: a bf16 send into an fp32-sized recv + # slot wedges the receiver forever). + recv_dtype = getattr(op.recv_shard_meta, "dtype", None) + if recv_dtype is not None and cloned.dtype != recv_dtype: + cloned = cloned.to(recv_dtype) + if not cloned.is_contiguous(): + cloned = cloned.contiguous() p2p_op = dist.P2POp( dist.isend if async_op else dist.send, - tensor_sliced.clone(), + cloned, recv_rank, group=weights_update_group, ) @@ -132,10 +199,9 @@ def update_weights_in_colocate_mode( for op in ops: recv_tensor = recv_parameters[op.recv_shard_meta.name] tensor_sliced = slice_tensor(recv_tensor, op, False) - if not tensor_sliced.is_contiguous(): - original_tensor = tensor_sliced - tensor_sliced = tensor_sliced.contiguous() - non_contiguous_tensor_pairs.append((original_tensor, tensor_sliced)) + tensor_sliced, copyback_pair = _prepare_p2p_recv_tensor(tensor_sliced) + if copyback_pair is not None: + recv_tensor_pairs.append(copyback_pair) p2p_op = dist.P2POp( dist.irecv if async_op else dist.recv, tensor_sliced, @@ -163,9 +229,12 @@ def update_weights_in_colocate_mode( msg = f"[{os.getpid()}] execute {total_send_ops} sends {total_recv_ops} recvs with recursive partition for {task_id}" hang_detector.submit(detect_hang, future, msg, [], timeout=60) - # Execute recursive partition transfer - # FIXME: batch_isend_irecv hang sometimes, seems `batch_isend_irecv` can't handle asymmetric p2p communication. - # so we use send/recv directly + # Recursive-partition butterfly with per-peer batch_isend_irecv (see + # _execute_ops_concurrent). The phase structure is symmetric and the + # data layer is verified fully consistent; the earlier deadlock was + # purely from submitting a whole half's ops in one batch. Issuing one + # batch per peer caps in-flight P2P at O(1) peer and stays + # deadlock-free. self.execute_recursive_partition_stream_transfer( transfer_rank, world_size, @@ -175,13 +244,14 @@ def update_weights_in_colocate_mode( rank_coordinate, step_id, ) - if non_contiguous_tensor_pairs: - with torch.no_grad(): - for original_tensor, recv_tensor in non_contiguous_tensor_pairs: - original_tensor.copy_(recv_tensor) - non_contiguous_tensor_pairs.clear() - del non_contiguous_tensor_pairs + device_util.synchronize() + if recv_tensor_pairs: + logger.info( + f"Syncing {len(recv_tensor_pairs)} non-contiguous recv buffers for {task_id}" + ) + _sync_p2p_recv_tensor_pairs(recv_tensor_pairs) + device_util.synchronize() future.set_result(True) duration = time.time() - start_time logger.info( @@ -306,48 +376,430 @@ def _execute_ops_concurrent(self, ops_dict, peer_ranks): Returns: Total number of ops executed """ - # Collect ops from all peers that have operations, along with their peer_rank - peer_ops_with_rank = [] - active_peer_ranks = [] - for peer_rank in peer_ranks: - if peer_rank in ops_dict: - peer_ops_with_rank.append((peer_rank, ops_dict[peer_rank])) - active_peer_ranks.append(peer_rank) - - if not peer_ops_with_rank: - return 0 - - # Allocate stream indices sequentially to active peer ranks for even distribution - # This ensures ranks are evenly distributed across available streams - peer_to_stream_idx = {} - for idx, peer_rank in enumerate(active_peer_ranks): - stream_idx = idx % len(self._stream_pool) - peer_to_stream_idx[peer_rank] = stream_idx - - # Find the maximum number of ops across all peers - max_ops = max(len(ops) for _, ops in peer_ops_with_rank) + # Per-peer batch_isend_irecv. Submitting a WHOLE half's ops in one + # batch_isend_irecv (the previous behaviour) deadlocks at 32-rank / + # PP4->PP1 asymmetric scale: round 0's other-half has 16 peers and + # thousands of ops, and flooding NCCL with that many concurrent P2P + # channels exhausts them so the GPU drain never completes (data layer + # verified fully symmetric — the failure is purely runtime concurrency + # scale). Instead we walk peers in ascending rank order and issue ONE + # batch_isend_irecv per peer, capping in-flight P2P at O(1) peer. + # + # This is deadlock-free because a recursive-partition phase is + # single-direction: in phase 1 every first-half rank only SENDS and + # every second-half rank only RECVS (phase 2 is the mirror). A + # (sender, receiver) pair's per-peer batch is matched by NCCL group + # on (src, dst, group); serializing peers cannot form a wait cycle + # since no rank both sends and receives within the same phase. (This + # is exactly why recursive partition's symmetric phases are safe and + # the circle-shift directed ring was not.) + # + # Both sides must walk peers in the SAME (ascending) order so the + # k-th batch on a sender pairs with the corresponding recv on the + # receiver. peer_ranks is already an ascending range here. + trace = os.environ.get("AWEX_P2P_TRACE", "").strip() in ("1", "true", "True") + my_rank = self.transfer_rank total_ops = 0 + for peer_rank in peer_ranks: + ops = ops_dict.get(peer_rank) + if not ops: + continue + p2p_ops = [p2p_op for _, p2p_op in ops] + if not p2p_ops: + continue + if trace: + logger.info( + f"[P2P-TRACE rank={my_rank}] peer={peer_rank} " + f"nops={len(p2p_ops)} -> batch_isend_irecv (pre-wait)" + ) + works = dist.batch_isend_irecv(p2p_ops) + for work in works: + work.wait() + if trace: + logger.info( + f"[P2P-TRACE rank={my_rank}] peer={peer_rank} " + f"work.wait returned (enqueued) -> synchronize (waiting peer)" + ) + # Force GPU completion before the next peer. work.wait() only + # blocks the CPU thread until the CUDA event records 'enqueued', + # not actual NCCL kernel completion; syncing per peer keeps + # in-flight P2P bounded to one peer and surfaces any hang at the + # offending peer rather than at a later chunk boundary. + if hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.synchronize() + if trace: + logger.info( + f"[P2P-TRACE rank={my_rank}] peer={peer_rank} " + f"synchronize done (drained peer)" + ) + total_ops += len(p2p_ops) + return total_ops - # Execute ops in round-robin fashion: one op from each peer per iteration - # This allows concurrent execution across multiple peers - work_handles = [] - for op_idx in range(max_ops): - for peer_rank, ops in peer_ops_with_rank: - if op_idx < len(ops): - _, p2p_op = ops[op_idx] - # Use the stream allocated to this peer to maintain ordering - stream_idx = peer_to_stream_idx[peer_rank] - stream = self._stream_pool[stream_idx] - with device_util.stream(stream): - result = p2p_op.op( - p2p_op.tensor, p2p_op.peer, group=p2p_op.group + def _run_chunked( + self, + *, + task_id, + step_id, + train_to_infer_device_mapping, + infer_to_train_device_mapping, + transfer_rank, + rank_coordinate, + world_size, + send_ops, + recv_ops, + recv_transfer_plan, + weights_update_group, + send_parameters, + recv_parameters, + async_op, + chunk_bytes, + ): + """Chunked send/recv for AWEX colocation. + + Cross-rank determinism: chunk N takes ops[N*step:(N+1)*step] from each + peer's per-peer ops list. plan_builder.build_local_transfer_plan + already sorts each peer's ops by (send_shard_meta.name, send_offset, + recv_offset) (transfer_plan.py:571), so rank A's send_ops[B] and rank + B's recv_ops[A] are aligned index-by-index. Same step_size on every + rank means matching send/recv pairs always land in the same chunk. + NCCL P2P pairs FIFO within (group, src, dst) so this preserves + protocol semantics. + + step_size is derived from chunk_bytes by sampling the per-op nbytes + from a representative op so that one chunk's clones approach but do + not exceed chunk_bytes. + + Local self-copy (tensors_to_copy) and self-recv-from-other-trains do + not consume clone memory and are emitted once up front. + """ + train_slice_context = {} + + local_train_rank = infer_to_train_device_mapping.get(transfer_rank) + tensors_to_copy = [] + local_self_recv_collected = [] + + send_per_peer = {} + for peer_rank, ops in send_ops.items(): + mapped_peer_rank = train_to_infer_device_mapping.get(peer_rank, peer_rank) + if mapped_peer_rank == transfer_rank: + for op in ops: + op_send_rank = getattr(op, "send_rank", None) + if ( + local_train_rank is not None + and op_send_rank is not None + and op_send_rank != local_train_rank + ): + local_self_recv_collected.append(op) + else: + if op.send_shard_meta.name not in send_parameters: + raise KeyError(op.send_shard_meta.name) + send_tensor = send_parameters[op.send_shard_meta.name] + tensor_sliced = slice_tensor( + send_tensor, op, True, slice_context=train_slice_context ) - if p2p_op.op is dist.isend or p2p_op.op is dist.irecv: - work_handles.append(result) - total_ops += 1 + tensors_to_copy.append(tensor_sliced) + else: + missing = [ + op.send_shard_meta.name + for op in ops + if op.send_shard_meta.name not in send_parameters + ] + if missing: + raise KeyError(missing[0]) + send_per_peer[mapped_peer_rank] = list(ops) + + recv_per_peer = {} + for send_rank, ops in recv_ops.items(): + recv_from_rank = train_to_infer_device_mapping[send_rank] + if recv_from_rank == transfer_rank: + continue + recv_per_peer[recv_from_rank] = list(ops) + + local_self_recv_built = [] + for op in local_self_recv_collected: + recv_buf = recv_parameters[op.recv_shard_meta.name] + recv_sliced = slice_tensor(recv_buf, op, False) + recv_sliced, copyback_pair = _prepare_p2p_recv_tensor(recv_sliced) + actual_send_rank = train_to_infer_device_mapping.get( + op.send_rank, op.send_rank + ) + p2p_op = dist.P2POp( + dist.irecv, + recv_sliced, + actual_send_rank, + group=weights_update_group, + ) + local_self_recv_built.append( + (actual_send_rank, op, p2p_op, copyback_pair) + ) - # Wait for all async operations to complete - for work in work_handles: - work.wait() + if len(tensors_to_copy) > 0: + send_rank_for_self = infer_to_train_device_mapping[transfer_rank] + execute_tensors_to_copy( + tensors_to_copy, + recv_transfer_plan.operations[send_rank_for_self], + recv_parameters, + f"tensor copy for {task_id}", + ) + else: + logger.info(f"No tensors to copy for {task_id}") + # The self-copy phase is done; release its densified slice cache so the + # staging copies don't stay alive for the rest of the chunked transfer. + tensors_to_copy = None + train_slice_context.clear() + + def _op_wire_bytes(op) -> int: + # Wire dtype is the receiver's dtype: send clones are cast to + # op.recv_shard_meta.dtype before posting (see the chunk loop + # below), so fp32 receivers cost 2x a bf16 train-side estimate. + dtype = getattr(op.recv_shard_meta, "dtype", None) or getattr( + op.send_shard_meta, "dtype", None + ) + elem_size = 2 + if dtype is not None: + try: + elem_size = int(dtype.itemsize) + except (AttributeError, TypeError): + try: + elem_size = torch.empty((), dtype=dtype).element_size() + except (TypeError, RuntimeError): + pass + overlap_shape = getattr(op, "overlap_shape", None) + if overlap_shape: + numel = 1 + for d in overlap_shape: + numel *= max(int(d), 1) + return max(numel * elem_size, 1) + shape = op.send_shard_meta.shape + numel = 1 + try: + for i, s in enumerate(op.train_slices or []): + dim_size = shape[i] if i < len(shape) else 1 + start = s.start if s.start is not None else 0 + stop = s.stop if s.stop is not None else dim_size + numel *= max(stop - start, 1) + except (TypeError, IndexError): + numel = 1 + for d in shape: + numel *= d + return max(numel * elem_size, 1) + + # AWEX_CHUNK_MB bounds the transient allocation of ONE chunk across + # ALL peers: chunk N clones send ops[N*step:(N+1)*step] for every + # send peer and stages non-contiguous recvs for every recv peer. + # The budget must therefore be divided by the AGGREGATE per-index + # wire bytes summed over peers, not a single peer's sample op — + # with 16 active peers the sample-op estimate would allocate ~16x + # the requested budget. + all_peer_ops = list(send_per_peer.values()) + list(recv_per_peer.values()) + max_len = max((len(ops) for ops in all_peer_ops), default=0) + max_index_bytes = 0 + for idx in range(max_len): + agg = sum( + _op_wire_bytes(ops[idx]) for ops in all_peer_ops if idx < len(ops) + ) + max_index_bytes = max(max_index_bytes, agg) - return total_ops + if max_index_bytes == 0: + step_size = 1 + else: + step_size = max(1, chunk_bytes // max_index_bytes) + logger.info( + f"[CHUNKED {task_id}] max_aggregate_bytes_per_index={max_index_bytes} " + f"chunk_bytes={chunk_bytes} step_size_local={step_size}" + ) + + env_force = os.environ.get("AWEX_CHUNK_OPS", "").strip() + if env_force: + try: + forced = max(1, int(env_force)) + step_size = forced + logger.info(f"[CHUNKED {task_id}] AWEX_CHUNK_OPS override={forced}") + except ValueError: + pass + else: + try: + if dist.is_initialized(): + t = torch.tensor( + [int(step_size)], + device=device_util.get_torch_device(), + dtype=torch.int64, + ) + dist.all_reduce( + t, op=dist.ReduceOp.MIN, group=weights_update_group + ) + new_step = int(t.item()) + if new_step != step_size: + logger.info( + f"[CHUNKED {task_id}] step_size aligned via all_reduce " + f"local={step_size} -> global_min={new_step}" + ) + step_size = max(1, new_step) + except Exception as e: + logger.warning( + f"[CHUNKED {task_id}] step_size all_reduce failed: {e}; " + f"using local={step_size} (risk of cross-rank chunk drift)" + ) + + max_send_len = max((len(v) for v in send_per_peer.values()), default=0) + max_recv_len = max((len(v) for v in recv_per_peer.values()), default=0) + n_chunks = max( + 1, + (max(max_send_len, max_recv_len) + step_size - 1) // step_size, + ) + # n_chunks must be globally consistent; otherwise ranks with fewer + # chunks exit the loop early and the others hang in batch_isend_irecv + # waiting for peers that already left. step_size is already MIN-reduced + # above, but n_chunks depends on per-rank max_send/recv lengths which + # diverge across ranks. Take MAX to ensure every rank runs the same + # number of chunk iterations (empty chunks are no-ops). + try: + if dist.is_initialized(): + t = torch.tensor( + [int(n_chunks)], + device=device_util.get_torch_device(), + dtype=torch.int64, + ) + dist.all_reduce( + t, op=dist.ReduceOp.MAX, group=weights_update_group + ) + new_n = int(t.item()) + if new_n != n_chunks: + logger.info( + f"[CHUNKED {task_id}] n_chunks aligned via all_reduce " + f"local={n_chunks} -> global_max={new_n}" + ) + n_chunks = new_n + except Exception as e: + logger.warning( + f"[CHUNKED {task_id}] n_chunks all_reduce failed: {e}; " + f"using local={n_chunks} (risk of cross-rank chunk drift / hang)" + ) + logger.info( + f"[CHUNKED {task_id}] n_chunks={n_chunks} step_size={step_size} " + f"max_send_per_peer={max_send_len} max_recv_per_peer={max_recv_len}" + ) + + total_clone_bytes = 0 + + for chunk_idx in range(n_chunks): + start = chunk_idx * step_size + end = start + step_size + + logger.warning( + f"[CHUNKED-DIAG {task_id}] chunk_idx={chunk_idx}/{n_chunks} ENTER " + f"slice=[{start},{end})" + ) + chunk_send_p2p_ops = {} + chunk_recv_p2p_ops = {} + chunk_recv_tensor_pairs = [] + chunk_clone_bytes = 0 + # Per-chunk slice cache: dedupes densified non-contiguous send + # slices WITHIN a chunk but is dropped at the chunk boundary, so + # dense staging copies do not accumulate across the whole + # transfer (which would defeat the AWEX_CHUNK_MB budget). + chunk_slice_context = {} + + for mapped_peer_rank, ops in send_per_peer.items(): + sub = ops[start:end] + if not sub: + continue + p2p_ops = [] + for op in sub: + send_tensor = send_parameters[op.send_shard_meta.name] + tensor_sliced = slice_tensor( + send_tensor, op, True, slice_context=chunk_slice_context + ) + recv_rank = train_to_infer_device_mapping.get( + op.recv_rank, op.recv_rank + ) + cloned = _clone_p2p_send_tensor(tensor_sliced) + # Wire-size parity: the receiver posts irecv with ITS shard + # dtype. Some parameters (e.g. mlp.gate.weight) are bf16 + # on the train side but fp32 on the sglang side; sending + # bf16 bytes into an fp32-sized recv leaves the receiver + # waiting forever (a deterministic mid-transfer deadlock). + # Cast the clone to the receiver's dtype. + recv_dtype = getattr(op.recv_shard_meta, "dtype", None) + if recv_dtype is not None and cloned.dtype != recv_dtype: + cloned = cloned.to(recv_dtype) + if not cloned.is_contiguous(): + cloned = cloned.contiguous() + p2p_op = dist.P2POp( + dist.isend if async_op else dist.send, + cloned, + recv_rank, + group=weights_update_group, + ) + p2p_ops.append((op, p2p_op)) + chunk_clone_bytes += cloned.numel() * cloned.element_size() + chunk_send_p2p_ops[mapped_peer_rank] = p2p_ops + + for recv_from_rank, ops in recv_per_peer.items(): + sub = ops[start:end] + if not sub: + continue + p2p_ops = [] + for op in sub: + recv_tensor = recv_parameters[op.recv_shard_meta.name] + tensor_sliced = slice_tensor(recv_tensor, op, False) + tensor_sliced, copyback_pair = _prepare_p2p_recv_tensor( + tensor_sliced + ) + if copyback_pair is not None: + chunk_recv_tensor_pairs.append(copyback_pair) + p2p_op = dist.P2POp( + dist.irecv if async_op else dist.recv, + tensor_sliced, + recv_from_rank, + group=weights_update_group, + ) + p2p_ops.append((op, p2p_op)) + chunk_recv_p2p_ops[recv_from_rank] = p2p_ops + + if chunk_idx == 0 and local_self_recv_built: + for actual_send_rank, op, p2p_op, copyback_pair in local_self_recv_built: + chunk_recv_p2p_ops.setdefault(actual_send_rank, []).append( + (op, p2p_op) + ) + if copyback_pair is not None: + chunk_recv_tensor_pairs.append(copyback_pair) + + self.execute_recursive_partition_stream_transfer( + transfer_rank, + world_size, + chunk_send_p2p_ops, + chunk_recv_p2p_ops, + weights_update_group, + rank_coordinate, + step_id, + ) + device_util.synchronize() + if chunk_recv_tensor_pairs: + logger.info( + f"[CHUNKED {task_id}] syncing {len(chunk_recv_tensor_pairs)} " + f"non-contiguous recv buffers for chunk {chunk_idx}" + ) + _sync_p2p_recv_tensor_pairs(chunk_recv_tensor_pairs) + device_util.synchronize() + logger.warning( + f"[CHUNKED-DIAG {task_id}] chunk_idx={chunk_idx}/{n_chunks} EXIT " + f"send_peers={len(chunk_send_p2p_ops)} recv_peers={len(chunk_recv_p2p_ops)} " + f"clone_mb={chunk_clone_bytes/1024/1024:.1f}" + ) + + chunk_send_p2p_ops = None + chunk_recv_p2p_ops = None + chunk_recv_tensor_pairs = None + chunk_slice_context = None + import gc as _gc + _gc.collect() + if hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.empty_cache() + + total_clone_bytes += chunk_clone_bytes + + logger.info( + f"CHUNKED transfer done {task_id}: chunks={n_chunks} step_size={step_size} " + f"total_clone_mb={total_clone_bytes / 1024 / 1024:.2f}" + ) diff --git a/awex/util/common.py b/awex/util/common.py index 8e8c089..a8cafd0 100644 --- a/awex/util/common.py +++ b/awex/util/common.py @@ -153,7 +153,7 @@ def compute_statistics(stage_history: dict, step_id: int, duration: float, stage history.append(duration) if len(history) > 10000: history.pop(0) - if step_id == 2: + if step_id == 2 and len(history) > 1: # first step contains init time history.pop(history.index(max(history))) num_updates = len(history)