From bbbd10ecb8fa6e9cfff88ee52bd5546119bf7610 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 15:45:47 +0800 Subject: [PATCH 1/8] fix: chunked per-peer P2P schedule and wire-size parity in colocate transfer Three fixes for colocate P2P at large asymmetric reshard scale (32 ranks, PP4 -> PP1), plus tracing: - Chunked transfer (AWEX_CHUNK_MB): split each peer's ops into size-bounded chunks with cross-rank aligned step_size/n_chunks (MIN/MAX all-reduced) so no rank exits the chunk loop early while its peers still wait. - Per-peer batch_isend_irecv with per-peer GPU drain: submitting a whole recursive-partition half's ops in one batch floods NCCL P2P channels and the GPU drain never completes. Walking peers in ascending order caps in-flight P2P at O(1) peer; phases are single-direction so serialization cannot form a wait cycle. - Wire-size parity: cast the send clone to the receiver's shard dtype. When the train-side shard is bf16 but the receiver posts an fp32-sized irecv (e.g. mlp.gate.weight), the byte counts differ and the receiver blocks forever. - AWEX_P2P_TRACE per-peer progress log for diagnosing peer-level stalls; execute_tensors_to_copy wrapped in no_grad. --- awex/transfer/nccl_comm.py | 1 + awex/transfer/nccl_stream_batch.py | 468 +++++++++++++++++++++++++---- 2 files changed, 413 insertions(+), 56 deletions(-) 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..c900fb8 100644 --- a/awex/transfer/nccl_stream_batch.py +++ b/awex/transfer/nccl_stream_batch.py @@ -35,7 +35,6 @@ logger = logging.getLogger(__name__) hang_detector = ThreadPoolExecutor(max_workers=1) - class NcclColocateStreamBatchTransport: MAX_STREAMS = 64 @@ -82,12 +81,39 @@ 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 = [] train_slice_context = {} - non_contiguous_tensor_pairs = [] # Process send operations for peer_rank, ops in send_ops.items(): @@ -113,9 +139,16 @@ def update_weights_in_colocate_mode( recv_rank = train_to_infer_device_mapping.get( op.recv_rank, op.recv_rank ) + cloned = tensor_sliced.clone() + # Wire-size parity with the receiver's dtype (see the + # chunked path / Problem 69: bf16 gate.weight into an fp32 + # 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) p2p_op = dist.P2POp( dist.isend if async_op else dist.send, - tensor_sliced.clone(), + cloned, recv_rank, group=weights_update_group, ) @@ -132,10 +165,6 @@ 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)) p2p_op = dist.P2POp( dist.irecv if async_op else dist.recv, tensor_sliced, @@ -163,9 +192,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,12 +207,7 @@ 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() future.set_result(True) duration = time.time() - start_time @@ -306,48 +333,377 @@ 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 + + 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 = {} - # 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 + 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) + 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)) - # 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}") - return total_ops + sample_op = None + for peer_rank in sorted(send_per_peer.keys()): + ops = send_per_peer[peer_rank] + if ops: + sample_op = ops[0] + break + if sample_op is None: + for peer_rank in sorted(recv_per_peer.keys()): + ops = recv_per_peer[peer_rank] + if ops: + sample_op = ops[0] + break + + if sample_op is None: + step_size = 1 + else: + shape = sample_op.send_shard_meta.shape + elem_size = 2 + try: + from awex.util.tensor_util import dtype_to_size as _dtype_size + elem_size = _dtype_size(sample_op.send_shard_meta.dtype) + except Exception: + pass + sliced_numel = 1 + try: + for s in (sample_op.train_slices or []): + span = s.stop - s.start if s.stop is not None else 0 + sliced_numel *= max(span, 1) + except Exception: + sliced_numel = 1 + for d in shape: + sliced_numel *= d + per_op_bytes = max(sliced_numel * elem_size, 1) + step_size = max(1, chunk_bytes // per_op_bytes) + logger.info( + f"[CHUNKED {task_id}] local sample shape={shape} per_op_bytes={per_op_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.current_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.current_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_clone_bytes = 0 + + 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=train_slice_context + ) + recv_rank = train_to_infer_device_mapping.get( + op.recv_rank, op.recv_rank + ) + cloned = tensor_sliced.clone() + # Wire-size parity: the receiver posts irecv with ITS shard + # dtype. 961 plan ops (mlp.gate.weight, 124 edges) 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 (deterministic chunk-7 deadlock, + # Problem 69). 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) + 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) + 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 in local_self_recv_built: + chunk_recv_p2p_ops.setdefault(actual_send_rank, []).append( + (op, p2p_op) + ) + + 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() + 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 + 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}" + ) From 88e7dbd9459375251c0d20a4ecb4dfea3551083b Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 15:45:50 +0800 Subject: [PATCH 2/8] fix: BailingMoe colocate adaptation and Lightning qkv TP declaration - Declare mcore Lightning attention qkv as TP_SHARDING when attn_tp_size > 1: convert_qkv_weight_along_tp_attention hands each rank only its shard, so declaring NO_SHARDING made the transfer plan cover infer tp_rank 0 only and left Lightning qkv on tp_rank > 0 all-zero. Regression test included. - Adapt the mcore converter for BailingMoeV2.5 colocate transfer and normalize SGLang's word_embeddings name to the canonical embed_tokens so both sides' transfer-plan key sets align. --- awex/converter/mcore_converter.py | 161 +++++++++++++++++---- awex/models/ling.py | 12 ++ awex/models/ling_linear.py | 12 +- awex/tests/test_attn_tp_sharding_cp_bug.py | 114 +++++++++++++++ 4 files changed, 270 insertions(+), 29 deletions(-) create mode 100644 awex/tests/test_attn_tp_sharding_cp_bug.py diff --git a/awex/converter/mcore_converter.py b/awex/converter/mcore_converter.py index d4c6da7..1110685 100644 --- a/awex/converter/mcore_converter.py +++ b/awex/converter/mcore_converter.py @@ -45,6 +45,52 @@ def _cfg_get(tf_config, key: str, default=None): return getattr(tf_config, key, default) +def pack_fused_qkv_a_proj_for_tp( + q_proj: torch.Tensor, + kv_proj: torch.Tensor, + infer_tp_size: int, +) -> torch.Tensor: + """Pack MLA q_a/kv_a full tensors in infer-TP local shard order. + + SGLang's fused MLA parameter is consumed per TP rank as: + [q_local_rank_i ; kv_local_rank_i] + + AWEX later applies generic TP chunking on dim 0, so the full writer-side + tensor must be laid out as: + [q_0 ; kv_0 ; q_1 ; kv_1 ; ...] + rather than [q_all ; kv_all]. + """ + if infer_tp_size <= 0: + raise ValueError(f"infer_tp_size must be positive, got {infer_tp_size}") + if q_proj.dim() != 2 or kv_proj.dim() != 2: + raise ValueError( + "Expected 2D MLA projection weights, got " + f"q_proj.dim={q_proj.dim()}, kv_proj.dim={kv_proj.dim()}" + ) + if q_proj.shape[1] != kv_proj.shape[1]: + raise ValueError( + "MLA q/kv projection input dims must match, got " + f"q_proj.shape={tuple(q_proj.shape)}, kv_proj.shape={tuple(kv_proj.shape)}" + ) + if q_proj.shape[0] % infer_tp_size != 0: + raise ValueError( + "MLA q projection rows must be divisible by infer_tp_size, got " + f"rows={q_proj.shape[0]}, infer_tp_size={infer_tp_size}" + ) + if kv_proj.shape[0] % infer_tp_size != 0: + raise ValueError( + "MLA kv projection rows must be divisible by infer_tp_size, got " + f"rows={kv_proj.shape[0]}, infer_tp_size={infer_tp_size}" + ) + + q_shards = q_proj.chunk(infer_tp_size, dim=0) + kv_shards = kv_proj.chunk(infer_tp_size, dim=0) + return torch.cat( + [torch.cat([q_shards[i], kv_shards[i]], dim=0) for i in range(infer_tp_size)], + dim=0, + ) + + def _normalize_pp_stage_layer_id_map( raw_map: Optional[Dict], ) -> Dict[Tuple[int, int], Dict[int, int]]: @@ -651,7 +697,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 +794,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 +825,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 +842,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,9 +887,25 @@ def _try_fuse_qkv_a_proj( if other_key not in layer_cache: return [] - fused_tensor = torch.cat( - [layer_cache["q_a_proj"], layer_cache["kv_a_proj"]], dim=0 - ) + # P97 fix: pack_fused_qkv_a_proj_for_tp lays the fused MLA a_proj out in + # infer-TP interleaved order [q_0;kv_0;q_1;kv_1;...], which is ONLY + # correct when AWEX subsequently TP-chunks this param back into per-rank + # [q_i;kv_i] shards (train attn_tp>1 path). When megatron attn has NO TP + # (attn_tp_size==1) the param is resolved as NO_SHARDING: each sglang TP + # rank receives the FULL tensor uncut, so it must match sglang's own + # load_weights layout `torch.cat([q_a, kv_a], dim=0)` exactly. The + # interleave otherwise scrambles the rows (HF-groundtruth rel=1.086, + # logp_diff blows up at step2). Emit the plain [q_a;kv_a] concat here. + if self.rank_info.attn_tp_size == 1: + fused_tensor = torch.cat( + [layer_cache["q_a_proj"], layer_cache["kv_a_proj"]], dim=0 + ) + else: + fused_tensor = pack_fused_qkv_a_proj_for_tp( + layer_cache["q_a_proj"], + layer_cache["kv_a_proj"], + self.infer_atten_tp_size, + ) del self.qkv_a_proj_cache[layer_number] return [("attention.fused_qkv_a_proj_with_mqa.weight", fused_tensor)] @@ -855,9 +941,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 +1105,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 +1143,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 = [] diff --git a/awex/models/ling.py b/awex/models/ling.py index 2c7549e..a0dffb8 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 (P73). + 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/tests/test_attn_tp_sharding_cp_bug.py b/awex/tests/test_attn_tp_sharding_cp_bug.py new file mode 100644 index 0000000..cf7ca29 --- /dev/null +++ b/awex/tests/test_attn_tp_sharding_cp_bug.py @@ -0,0 +1,114 @@ +# Copyright (c) Ant Group. Licensed under the Apache License, Version 2.0. +""" +Regression test for the AWEX P73-class bug on **train attn_tp == 1 + infer TP > 1**. + +Background (root cause, 2026-06-17 investigation): + 共卡 SWE 实验 (allocation actor attn=d2p4c8, 即 attn_tp_size==1 纯 CP8) 在 step2 + 推理崩溃 (reject 98%), 而 zjw / gsm8k 用 attn=d1p4t4c2 (attn_tp_size==4) 正常。 + 唯一区分变量收敛到 actor attn 是否开 TP。 + + awex `BailingMoeShardingStrategy.get_sharding_strategy` 对 query_key_value 的注释 + (ling.py) 已记录同款失败模式 (P73): + > Declaring NO_SHARDING ... the transfer plan generated ops for infer + > tp_rank 0 only and Lightning qkv on tp_rank > 0 stayed all-zero. + P73 的修复只覆盖 `attn_tp_size > 1` 分支 (返回 TP_SHARDING, 对齐 infer TP)； + 当 train `attn_tp_size == 1` 时仍回落到 NO_SHARDING —— 但推理侧是 TP>1 (sglang + d16t4)，于是 transfer plan 同样只覆盖 infer tp_rank 0、tp_rank>0 权重保持全零， + step2 推理生成崩溃。 + +这个测试把上述「触发条件」固化为回归用例： + - attn_tp==4 (正常拓扑) -> query_key_value 走 TP_SHARDING (对齐 infer TP) + - attn_tp==1 (崩溃拓扑) -> query_key_value 回落 NO_SHARDING (P73 同款 bug 触发条件) + +注意：本测试在 sharding-strategy 层坐实「触发条件」(无 TP -> NO_SHARDING 声明)。 +完整的「transfer plan 漏给 infer tp_rank>0 生成 op」需要 meta-resolver + plan 生成 +层的端到端测试 (TODO)，并已由 lr=0 运行时实验 (step2 behave_imp_weight 偏离 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: + """构造 train(mcore)侧 RankInfo。attn_tp_size==1 + cp_size>1 = 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(): + """正常拓扑 d1p4t4c2 / d2p4t4c2: attn_tp=4 -> TP_SHARDING, 和 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(): + """崩溃拓扑 d2p4c8: attn_tp=1(纯CP8) -> NO_SHARDING。 + + 这是 P73 同款 bug 的触发条件: train 声明完整副本(NO_SHARDING), 但推理侧 sglang + 是 TP>1, transfer plan 只给 infer tp_rank 0 生成 op, tp_rank>0 权重全零 -> 崩。 + P73 的修复仅覆盖 attn_tp>1 分支, 未覆盖此处 attn_tp==1 + infer_tp>1 组合。 + """ + strat = _make_strategy(attn_tp_size=1, cp_size=8) + sharding_type, _dim, num_shards = strat.get_sharding_strategy(QKV_PARAM) + # 当前 awex 行为(bug 触发条件): 回落 NO_SHARDING。 + assert sharding_type == ShardingType.NO_SHARDING + assert num_shards == 1 + + +def test_cp_size_does_not_change_qkv_strategy(): + """对照: 固定 attn_tp=1, cp_size 从 1 变到 8, query_key_value 的 sharding 决策不变。 + + 证明 awex 的 attn sharding 决策完全不感知 CP 维度——CP8 的 8 份冗余副本 + 在 strategy 层没有任何特殊处理, 全部依赖下游收集/传输层, 而该层在 + attn_tp==1 + infer_tp>1 时漏传(P73 同款)。 + """ + 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 不影响 attn sharding 决策(盲区) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From 562e83d805c58a32abf8407a4afad3f78b49e59c Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 15:45:52 +0800 Subject: [PATCH 3/8] fix: force IPC close before signalling weights_update_finished cudaIpcCloseMemHandle runs in the tensor deleter, which only fires when 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 hits an illegal memory access at a drifting location. gc.collect + synchronize before the signal. --- awex/reader/nccl_reader.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/awex/reader/nccl_reader.py b/awex/reader/nccl_reader.py index e576b60..7be11d6 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 at a drifting location + # (Problem 71: post-transfer IMA in reload / resume / first read). + # Force the close to complete BEFORE signalling the train side. self.deserialized_weights = None + gc.collect() + torch.cuda.synchronize() duration = time.time() - start_time compute_statistics( self._history_update_weights_time, From e2dd320546c237d30d9dd7c943c23f443ac122ee Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 15:45:54 +0800 Subject: [PATCH 4/8] fix: guard update-history compare against short history Skip the step-2 consistency compare when fewer than two history entries exist, instead of crashing on single-entry histories. --- awex/util/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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) From 024fb0902f3a7793528fba7834d06a5df2da818a Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 08:29:01 +0800 Subject: [PATCH 5/8] fix(transfer): densify non-contiguous tensors for P2P send/recv torch.distributed P2P ops require dense buffers; KDA/MLA converter outputs are frequently non-contiguous views. Clone sends into contiguous buffers and stage non-contiguous receives through a dense buffer with an explicit copy-back after wait. Co-Authored-By: Claude Fable 5 --- awex/tests/test_nccl_stream_batch.py | 62 ++++++++++++++++++++++++ awex/transfer/nccl_stream_batch.py | 72 ++++++++++++++++++++++++++-- 2 files changed, 130 insertions(+), 4 deletions(-) create mode 100644 awex/tests/test_nccl_stream_batch.py 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_stream_batch.py b/awex/transfer/nccl_stream_batch.py index c900fb8..b6951f8 100644 --- a/awex/transfer/nccl_stream_batch.py +++ b/awex/transfer/nccl_stream_batch.py @@ -35,6 +35,37 @@ logger = logging.getLogger(__name__) 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 @@ -113,6 +144,7 @@ def update_weights_in_colocate_mode( 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 = {} # Process send operations @@ -139,13 +171,15 @@ def update_weights_in_colocate_mode( recv_rank = train_to_infer_device_mapping.get( op.recv_rank, op.recv_rank ) - cloned = tensor_sliced.clone() + cloned = _clone_p2p_send_tensor(tensor_sliced) # Wire-size parity with the receiver's dtype (see the # chunked path / Problem 69: bf16 gate.weight into an fp32 # 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, cloned, @@ -165,6 +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) + 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, @@ -209,6 +246,12 @@ def update_weights_in_colocate_mode( ) 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( @@ -476,6 +519,7 @@ def _run_chunked( 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 ) @@ -485,7 +529,9 @@ def _run_chunked( actual_send_rank, group=weights_update_group, ) - local_self_recv_built.append((actual_send_rank, op, p2p_op)) + local_self_recv_built.append( + (actual_send_rank, op, p2p_op, copyback_pair) + ) if len(tensors_to_copy) > 0: send_rank_for_self = infer_to_train_device_mapping[transfer_rank] @@ -620,6 +666,7 @@ def _run_chunked( ) chunk_send_p2p_ops = {} chunk_recv_p2p_ops = {} + chunk_recv_tensor_pairs = [] chunk_clone_bytes = 0 for mapped_peer_rank, ops in send_per_peer.items(): @@ -635,7 +682,7 @@ def _run_chunked( recv_rank = train_to_infer_device_mapping.get( op.recv_rank, op.recv_rank ) - cloned = tensor_sliced.clone() + cloned = _clone_p2p_send_tensor(tensor_sliced) # Wire-size parity: the receiver posts irecv with ITS shard # dtype. 961 plan ops (mlp.gate.weight, 124 edges) are bf16 # on the train side but fp32 on the sglang side; sending @@ -645,6 +692,8 @@ def _run_chunked( 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, @@ -663,6 +712,11 @@ def _run_chunked( 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, @@ -673,10 +727,12 @@ def _run_chunked( 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 in 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, @@ -688,6 +744,13 @@ def _run_chunked( 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)} " @@ -696,6 +759,7 @@ def _run_chunked( chunk_send_p2p_ops = None chunk_recv_p2p_ops = None + chunk_recv_tensor_pairs = None import gc as _gc _gc.collect() if hasattr(torch, "cuda") and torch.cuda.is_available(): From 8425cb06a4aecf01101f1fd8a896247e01c69934 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 16:54:42 +0800 Subject: [PATCH 6/8] docs: translate test comments to English and drop internal issue tags Address review: rewrite test_attn_tp_sharding_cp_bug.py docs/comments in English and replace internal issue references in code comments with self-contained failure-mode descriptions. --- awex/models/ling.py | 2 +- awex/reader/nccl_reader.py | 6 +- awex/tests/test_attn_tp_sharding_cp_bug.py | 86 +++++++++++++--------- awex/transfer/nccl_stream_batch.py | 22 +++--- 4 files changed, 68 insertions(+), 48 deletions(-) diff --git a/awex/models/ling.py b/awex/models/ling.py index a0dffb8..c147089 100644 --- a/awex/models/ling.py +++ b/awex/models/ling.py @@ -45,7 +45,7 @@ def get_sharding_strategy(self, parameter_name, **kwargs): # 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 (P73). + # all-zero. attn_tp_size = self.rank_info.attn_tp_size if attn_tp_size > 1: return ShardingType.TP_SHARDING, 0, attn_tp_size diff --git a/awex/reader/nccl_reader.py b/awex/reader/nccl_reader.py index 7be11d6..c548bd0 100644 --- a/awex/reader/nccl_reader.py +++ b/awex/reader/nccl_reader.py @@ -468,12 +468,12 @@ def _update_weights_in_colocate_mode(self, step_id, **kwargs): # 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 at a drifting location - # (Problem 71: post-transfer IMA in reload / resume / first read). + # 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() - torch.cuda.synchronize() + device_util.synchronize() duration = time.time() - start_time compute_statistics( self._history_update_weights_time, diff --git a/awex/tests/test_attn_tp_sharding_cp_bug.py b/awex/tests/test_attn_tp_sharding_cp_bug.py index cf7ca29..a83eea6 100644 --- a/awex/tests/test_attn_tp_sharding_cp_bug.py +++ b/awex/tests/test_attn_tp_sharding_cp_bug.py @@ -1,28 +1,38 @@ # Copyright (c) Ant Group. Licensed under the Apache License, Version 2.0. """ -Regression test for the AWEX P73-class bug on **train attn_tp == 1 + infer TP > 1**. - -Background (root cause, 2026-06-17 investigation): - 共卡 SWE 实验 (allocation actor attn=d2p4c8, 即 attn_tp_size==1 纯 CP8) 在 step2 - 推理崩溃 (reject 98%), 而 zjw / gsm8k 用 attn=d1p4t4c2 (attn_tp_size==4) 正常。 - 唯一区分变量收敛到 actor attn 是否开 TP。 - - awex `BailingMoeShardingStrategy.get_sharding_strategy` 对 query_key_value 的注释 - (ling.py) 已记录同款失败模式 (P73): +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. - P73 的修复只覆盖 `attn_tp_size > 1` 分支 (返回 TP_SHARDING, 对齐 infer TP)； - 当 train `attn_tp_size == 1` 时仍回落到 NO_SHARDING —— 但推理侧是 TP>1 (sglang - d16t4)，于是 transfer plan 同样只覆盖 infer tp_rank 0、tp_rank>0 权重保持全零， - step2 推理生成崩溃。 - -这个测试把上述「触发条件」固化为回归用例： - - attn_tp==4 (正常拓扑) -> query_key_value 走 TP_SHARDING (对齐 infer TP) - - attn_tp==1 (崩溃拓扑) -> query_key_value 回落 NO_SHARDING (P73 同款 bug 触发条件) - -注意：本测试在 sharding-strategy 层坐实「触发条件」(无 TP -> NO_SHARDING 声明)。 -完整的「transfer plan 漏给 infer tp_rank>0 生成 op」需要 meta-resolver + plan 生成 -层的端到端测试 (TODO)，并已由 lr=0 运行时实验 (step2 behave_imp_weight 偏离 1) 互证。 + 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 @@ -32,7 +42,11 @@ def _make_rank_info(attn_tp_size: int, cp_size: int) -> RankInfo: - """构造 train(mcore)侧 RankInfo。attn_tp_size==1 + cp_size>1 = d2p4c8 崩溃拓扑。""" + """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, @@ -77,7 +91,7 @@ def _make_strategy(attn_tp_size: int, cp_size: int): def test_qkv_sharding_with_attn_tp4_is_tp_sharded(): - """正常拓扑 d1p4t4c2 / d2p4t4c2: attn_tp=4 -> TP_SHARDING, 和 infer TP 对齐, 传对。""" + """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 @@ -85,29 +99,33 @@ def test_qkv_sharding_with_attn_tp4_is_tp_sharded(): def test_qkv_sharding_with_attn_tp1_falls_back_to_no_sharding(): - """崩溃拓扑 d2p4c8: attn_tp=1(纯CP8) -> NO_SHARDING。 - - 这是 P73 同款 bug 的触发条件: train 声明完整副本(NO_SHARDING), 但推理侧 sglang - 是 TP>1, transfer plan 只给 infer tp_rank 0 生成 op, tp_rank>0 权重全零 -> 崩。 - P73 的修复仅覆盖 attn_tp>1 分支, 未覆盖此处 attn_tp==1 + infer_tp>1 组合。 + """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) - # 当前 awex 行为(bug 触发条件): 回落 NO_SHARDING。 + # 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(): - """对照: 固定 attn_tp=1, cp_size 从 1 变到 8, query_key_value 的 sharding 决策不变。 + """Control: with attn_tp=1 fixed, cp_size 1 -> 8 does not change the decision. - 证明 awex 的 attn sharding 决策完全不感知 CP 维度——CP8 的 8 份冗余副本 - 在 strategy 层没有任何特殊处理, 全部依赖下游收集/传输层, 而该层在 - attn_tp==1 + infer_tp>1 时漏传(P73 同款)。 + 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 不影响 attn sharding 决策(盲区) + assert s_cp1 == s_cp8 # cp_size does not affect the attn sharding decision if __name__ == "__main__": diff --git a/awex/transfer/nccl_stream_batch.py b/awex/transfer/nccl_stream_batch.py index b6951f8..a5f3f3e 100644 --- a/awex/transfer/nccl_stream_batch.py +++ b/awex/transfer/nccl_stream_batch.py @@ -173,8 +173,8 @@ def update_weights_in_colocate_mode( ) cloned = _clone_p2p_send_tensor(tensor_sliced) # Wire-size parity with the receiver's dtype (see the - # chunked path / Problem 69: bf16 gate.weight into an fp32 - # recv slot wedges the receiver forever). + # 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) @@ -569,9 +569,11 @@ def _run_chunked( pass sliced_numel = 1 try: - for s in (sample_op.train_slices or []): - span = s.stop - s.start if s.stop is not None else 0 - sliced_numel *= max(span, 1) + for i, s in enumerate(sample_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 + sliced_numel *= max(stop - start, 1) except Exception: sliced_numel = 1 for d in shape: @@ -596,7 +598,7 @@ def _run_chunked( if dist.is_initialized(): t = torch.tensor( [int(step_size)], - device=device_util.current_device(), + device=device_util.get_torch_device(), dtype=torch.int64, ) dist.all_reduce( @@ -631,7 +633,7 @@ def _run_chunked( if dist.is_initialized(): t = torch.tensor( [int(n_chunks)], - device=device_util.current_device(), + device=device_util.get_torch_device(), dtype=torch.int64, ) dist.all_reduce( @@ -684,11 +686,11 @@ def _run_chunked( ) cloned = _clone_p2p_send_tensor(tensor_sliced) # Wire-size parity: the receiver posts irecv with ITS shard - # dtype. 961 plan ops (mlp.gate.weight, 124 edges) are bf16 + # 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 (deterministic chunk-7 deadlock, - # Problem 69). Cast the clone to the receiver's dtype. + # 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) From 6a897deb483780f7fef77ad65b2e7c7cb104af98 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Wed, 22 Jul 2026 16:39:06 +0800 Subject: [PATCH 7/8] fix(converter): keep canonical MLA fused concat and consistent qkv layout detection - emit the plain [q_a;kv_a] concat for fused_qkv_a_proj_with_mqa unconditionally: the param is declared NO_SHARDING, so every SGLang TP rank receives the full tensor and any infer-TP interleave scrambles rows; drop the now-unused pack_fused_qkv_a_proj_for_tp helper - check the compact GQA/MHA layout before the replicated one in transform_mcore_qkv_bias, mirroring transform_mcore_qkv_weight, so MHA (where both sizes coincide) resolves weight and bias identically - treat embed_tokens as an embedding in the sharding fallback so the canonical HF name keeps replicated-embedding semantics instead of falling through to generic TP_SHARDING --- awex/converter/mcore_converter.py | 115 +++++++++--------------------- awex/sharding/param_sharding.py | 6 +- 2 files changed, 39 insertions(+), 82 deletions(-) diff --git a/awex/converter/mcore_converter.py b/awex/converter/mcore_converter.py index 1110685..ccdfec9 100644 --- a/awex/converter/mcore_converter.py +++ b/awex/converter/mcore_converter.py @@ -45,52 +45,6 @@ def _cfg_get(tf_config, key: str, default=None): return getattr(tf_config, key, default) -def pack_fused_qkv_a_proj_for_tp( - q_proj: torch.Tensor, - kv_proj: torch.Tensor, - infer_tp_size: int, -) -> torch.Tensor: - """Pack MLA q_a/kv_a full tensors in infer-TP local shard order. - - SGLang's fused MLA parameter is consumed per TP rank as: - [q_local_rank_i ; kv_local_rank_i] - - AWEX later applies generic TP chunking on dim 0, so the full writer-side - tensor must be laid out as: - [q_0 ; kv_0 ; q_1 ; kv_1 ; ...] - rather than [q_all ; kv_all]. - """ - if infer_tp_size <= 0: - raise ValueError(f"infer_tp_size must be positive, got {infer_tp_size}") - if q_proj.dim() != 2 or kv_proj.dim() != 2: - raise ValueError( - "Expected 2D MLA projection weights, got " - f"q_proj.dim={q_proj.dim()}, kv_proj.dim={kv_proj.dim()}" - ) - if q_proj.shape[1] != kv_proj.shape[1]: - raise ValueError( - "MLA q/kv projection input dims must match, got " - f"q_proj.shape={tuple(q_proj.shape)}, kv_proj.shape={tuple(kv_proj.shape)}" - ) - if q_proj.shape[0] % infer_tp_size != 0: - raise ValueError( - "MLA q projection rows must be divisible by infer_tp_size, got " - f"rows={q_proj.shape[0]}, infer_tp_size={infer_tp_size}" - ) - if kv_proj.shape[0] % infer_tp_size != 0: - raise ValueError( - "MLA kv projection rows must be divisible by infer_tp_size, got " - f"rows={kv_proj.shape[0]}, infer_tp_size={infer_tp_size}" - ) - - q_shards = q_proj.chunk(infer_tp_size, dim=0) - kv_shards = kv_proj.chunk(infer_tp_size, dim=0) - return torch.cat( - [torch.cat([q_shards[i], kv_shards[i]], dim=0) for i in range(infer_tp_size)], - dim=0, - ) - - def _normalize_pp_stage_layer_id_map( raw_map: Optional[Dict], ) -> Dict[Tuple[int, int], Dict[int, int]]: @@ -887,25 +841,20 @@ def _try_fuse_qkv_a_proj( if other_key not in layer_cache: return [] - # P97 fix: pack_fused_qkv_a_proj_for_tp lays the fused MLA a_proj out in - # infer-TP interleaved order [q_0;kv_0;q_1;kv_1;...], which is ONLY - # correct when AWEX subsequently TP-chunks this param back into per-rank - # [q_i;kv_i] shards (train attn_tp>1 path). When megatron attn has NO TP - # (attn_tp_size==1) the param is resolved as NO_SHARDING: each sglang TP - # rank receives the FULL tensor uncut, so it must match sglang's own - # load_weights layout `torch.cat([q_a, kv_a], dim=0)` exactly. The - # interleave otherwise scrambles the rows (HF-groundtruth rel=1.086, - # logp_diff blows up at step2). Emit the plain [q_a;kv_a] concat here. - if self.rank_info.attn_tp_size == 1: - fused_tensor = torch.cat( - [layer_cache["q_a_proj"], layer_cache["kv_a_proj"]], dim=0 - ) - else: - fused_tensor = pack_fused_qkv_a_proj_for_tp( - layer_cache["q_a_proj"], - layer_cache["kv_a_proj"], - self.infer_atten_tp_size, - ) + # 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 + ) del self.qkv_a_proj_cache[layer_number] return [("attention.fused_qkv_a_proj_with_mqa.weight", fused_tensor)] @@ -1248,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) @@ -1270,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/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) From 5109a757188189d74f0a48e3405cea6627b45b23 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Wed, 22 Jul 2026 16:39:07 +0800 Subject: [PATCH 8/8] fix(transfer): make AWEX_CHUNK_MB a real bound on per-chunk allocation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - derive step_size from the aggregate wire bytes of one op index summed across ALL send/recv peers (using the receiver dtype, which the clones are cast to), instead of sampling a single peer's op — with N active peers the old estimate allocated up to N x the requested budget - scope the densified-send slice cache to each chunk and release the self-copy cache right after the copy phase, so staging buffers no longer accumulate across the whole transfer - drop the try/except around the removed dtype_to_size import that silently fell back to 2-byte elements for fp32 sends --- awex/transfer/nccl_stream_batch.py | 96 ++++++++++++++++++++---------- 1 file changed, 63 insertions(+), 33 deletions(-) diff --git a/awex/transfer/nccl_stream_batch.py b/awex/transfer/nccl_stream_batch.py index a5f3f3e..82e4eeb 100644 --- a/awex/transfer/nccl_stream_batch.py +++ b/awex/transfer/nccl_stream_batch.py @@ -543,45 +543,69 @@ def _run_chunked( ) else: logger.info(f"No tensors to copy for {task_id}") - - sample_op = None - for peer_rank in sorted(send_per_peer.keys()): - ops = send_per_peer[peer_rank] - if ops: - sample_op = ops[0] - break - if sample_op is None: - for peer_rank in sorted(recv_per_peer.keys()): - ops = recv_per_peer[peer_rank] - if ops: - sample_op = ops[0] - break - - if sample_op is None: - step_size = 1 - else: - shape = sample_op.send_shard_meta.shape + # 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: - from awex.util.tensor_util import dtype_to_size as _dtype_size - elem_size = _dtype_size(sample_op.send_shard_meta.dtype) - except Exception: - pass - sliced_numel = 1 - try: - for i, s in enumerate(sample_op.train_slices or []): + 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 - sliced_numel *= max(stop - start, 1) - except Exception: - sliced_numel = 1 + numel *= max(stop - start, 1) + except (TypeError, IndexError): + numel = 1 for d in shape: - sliced_numel *= d - per_op_bytes = max(sliced_numel * elem_size, 1) - step_size = max(1, chunk_bytes // per_op_bytes) + 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) + + if max_index_bytes == 0: + step_size = 1 + else: + step_size = max(1, chunk_bytes // max_index_bytes) logger.info( - f"[CHUNKED {task_id}] local sample shape={shape} per_op_bytes={per_op_bytes} " + f"[CHUNKED {task_id}] max_aggregate_bytes_per_index={max_index_bytes} " f"chunk_bytes={chunk_bytes} step_size_local={step_size}" ) @@ -670,6 +694,11 @@ def _run_chunked( 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] @@ -679,7 +708,7 @@ def _run_chunked( for op in sub: send_tensor = send_parameters[op.send_shard_meta.name] tensor_sliced = slice_tensor( - send_tensor, op, True, slice_context=train_slice_context + send_tensor, op, True, slice_context=chunk_slice_context ) recv_rank = train_to_infer_device_mapping.get( op.recv_rank, op.recv_rank @@ -762,6 +791,7 @@ def _run_chunked( 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():