From 7571d210723fe80099a4620932ce07e1479f5b34 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Mon, 14 Sep 2026 21:14:27 +0800 Subject: [PATCH 1/6] init --- src/mcore_bridge/bridge/gpt_bridge.py | 87 +++++++++++++++++++++++- src/mcore_bridge/model/gpts/qwen4_exp.py | 15 ++++ src/mcore_bridge/model/modules/ple.py | 66 ++++++++++++------ 3 files changed, 148 insertions(+), 20 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 41caf7cb..d2c2c43c 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -28,6 +28,9 @@ class GPTBridge: fp8_block_size = 128 + # Bound the per-collective GPU buffer when a full gathered tensor is + # streamed to CPU (checkpoint save / CPU-offloaded weight sync). + export_chunk_bytes = 256 << 20 hf_layers_prefix = 'model.layers' hf_mtp_prefix = 'model.layers' hf_embed_key = 'model.embed_tokens.weight' @@ -336,11 +339,57 @@ def _set_module(self, mg_module, hf_state_dict, hf_prefix: str, to_mcore: bool): hf_state_dict[k] = v.to(self._target_device) return self._add_prefix(hf_state_dict, hf_prefix) + def _stream_to_cpu(self) -> bool: + """Whether large gathers/broadcasts should be chunked straight into host + memory. Enabled when the exported tensor is requested on CPU (checkpoint + save / CPU-offloaded weight sync): materializing the full gathered tensor + on GPU can OOM while the training state is still resident. The decision + only depends on properties identical across the group, so the collective + sequence stays in sync on every rank. Every participating rank still + assembles the full CPU result (matching the pre-chunking semantics); only + the GPU-resident full-size buffer is avoided.""" + return self._target_device == 'cpu' + + def _chunk_rows(self, shape, elem_size: int) -> int: + """Rows (along dim0) per broadcast/gather chunk, ~export_chunk_bytes.""" + inner = 1 + for s in shape[1:]: + inner *= s + rows = max(1, self.export_chunk_bytes // max(1, inner * elem_size)) + return min(rows, shape[0]) + + def _chunked_all_gather_tp(self, tensor, tp_dim: int, tp_group, tp_size: int): + """All-gather `tensor` along tp_dim and assemble the result in host memory + chunk by chunk, so the full gathered tensor never exists on GPU at once. + The chunk count derives from the local shape, which is identical on every + rank of the group, so the collective sequence stays in sync.""" + dim_size = tensor.shape[tp_dim] + chunk_bytes = max(1, self.export_chunk_bytes // tp_size) + inner_rows = max(1, chunk_bytes // max(1, tensor.numel() // dim_size * tensor.element_size())) + out_shape = list(tensor.shape) + out_shape[tp_dim] = dim_size * tp_size + output = torch.empty(out_shape, dtype=tensor.dtype, device='cpu') + for start in range(0, dim_size, inner_rows): + end = min(dim_size, start + inner_rows) + local = tensor.narrow(tp_dim, start, end - start).contiguous() + gathered = [torch.empty_like(local) for _ in range(tp_size)] + dist.all_gather(gathered, local, group=tp_group) + del local + for j in range(tp_size): + dst = tuple( + slice(j * dim_size + start, j * dim_size + end) if ax == tp_dim else slice(None) + for ax in range(tensor.ndim)) + output[dst] = gathered[j].cpu() + del gathered + return output + def _all_gather_tp(self, tensor, tp_dim, is_expert): tensor = None if tensor is None else tensor.to('cuda') tp_size = self.etp_size if is_expert else self.tp_size tp_group = self.etp_group if is_expert else self.tp_group if tensor is not None and tp_dim is not None and tp_size > 1: + if self._stream_to_cpu() and tensor.numel() * tensor.element_size() > self.export_chunk_bytes: + return self._chunked_all_gather_tp(tensor, tp_dim, tp_group, tp_size) if tp_dim == 0: # save memory tensor_shape = list(tensor.shape) @@ -363,6 +412,33 @@ def _all_gather_tp(self, tensor, tp_dim, is_expert): del output return tensor + def _chunked_broadcast_pp(self, tensor, shape, dtype, src_rank: int, pp_group): + """Chunked pp/ep-pp broadcast (the pp counterpart of _chunked_all_gather_tp): + stream the tensor chunk by chunk instead of materializing the full buffer + on GPU. On the holder rank `tensor` carries the data and is returned + as-is; receivers pass `shape`/`dtype` from the already-broadcast meta and + assemble the result in host memory (ranks that do not keep the export + still join every collective but skip the host assembly). Both sides + derive the same chunk count from the meta shape.""" + rows = self._chunk_rows(shape, torch.tensor([], dtype=dtype).element_size()) + if tensor is not None: + for start in range(0, shape[0], rows): + end = min(shape[0], start + rows) + send = tensor[start:end] + if not send.is_cuda or send.dtype != dtype or not send.is_contiguous(): + send = send.to(device='cuda', dtype=dtype).contiguous() + dist.broadcast(send, src=src_rank, group=pp_group) + return tensor + output = torch.empty(shape, dtype=dtype, device='cpu') + buf = None + for start in range(0, shape[0], rows): + end = min(shape[0], start + rows) + if buf is None or buf.shape[0] != end - start: + buf = torch.empty([end - start] + list(shape[1:]), device='cuda', dtype=dtype) + dist.broadcast(buf, src=src_rank, group=pp_group) + output[start:end] = buf.cpu() + return output + def _broadcast_ep_pp(self, tensor, is_expert): pp_group = self.ep_pp_group if is_expert else self.pp_group pp_size = self.ep_pp_size if is_expert else self.pp_size @@ -379,6 +455,10 @@ def _broadcast_ep_pp(self, tensor, is_expert): dist.broadcast(meta_data, src=src_rank, group=pp_group) shape = meta_data[1:1 + meta_data[0]].tolist() dtype = dtype_mapping[meta_data[-1].item()] + numel = math.prod(shape) if shape else 1 + if self._stream_to_cpu() and numel * torch.empty( + (), dtype=dtype).element_size() > (self.export_chunk_bytes) and len(shape) > 0: + return self._chunked_broadcast_pp(None, shape, dtype, src_rank, pp_group) tensor = torch.empty(shape, device='cuda', dtype=dtype) dist.broadcast(tensor, src=src_rank, group=pp_group) else: @@ -386,6 +466,9 @@ def _broadcast_ep_pp(self, tensor, is_expert): meta_data[1:1 + tensor.ndim] = torch.tensor(tensor.shape, dtype=torch.int64, device='cuda') meta_data[-1] = dtype_mapping_r[tensor.dtype] dist.broadcast(meta_data, src=src_rank, group=pp_group) + if self._stream_to_cpu() and tensor.numel() * tensor.element_size() > ( + self.export_chunk_bytes) and tensor.ndim > 0: + return self._chunked_broadcast_pp(tensor, list(tensor.shape), tensor.dtype, src_rank, pp_group) dist.broadcast(tensor, src=src_rank, group=pp_group) return tensor @@ -766,7 +849,9 @@ def _get_hf_experts_attr(self, is_mtp: bool = False): if (self._is_saving and not is_mtp and not self.config.fp8_param and not self._peft_format and self.model_type == 'qwen3_5_moe'): return True, True - if self.model_type in {'glm4v_moe', 'kimi_vl', 'qwen3_omni_moe', 'qwen3_5_moe'} or self.llm_model_type in { + if self.model_type in { + 'glm4v_moe', 'glm5_next', 'kimi_vl', 'qwen3_omni_moe', 'qwen3_5_moe' + } or self.llm_model_type in { 'qwen2_moe', 'qwen3_moe', 'deepseek_v2', 'deepseek_v3', 'kimi_k2', 'dots1', 'ernie4_5_moe', 'glm4_moe', 'glm4_moe_lite', 'minimax_m2', 'olmoe', 'qwen3_next', 'glm_moe_dsa', 'deepseek_v32', 'deepseek_v4' }: diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 03b47681..e8c08582 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -23,6 +23,7 @@ from ..modules import (QSA_SPARSE_KERNEL_ENV, GatedDeltaNet, QSAIndexer, QSASparseCoreAttention, Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, TransformerBlock, TransformerLayer, qsa_sparse_supported, use_qsa_sparse_kernel) +from ..modules.ple import Qwen4ExpTextNGramEmbedding from ..register import ModelLoader from .qwen3_next import Qwen3NextBridge, Qwen3NextRMSNorm, Qwen3NextSelfAttention @@ -391,6 +392,20 @@ def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): ple.ple_embedding.fill_table_from_hf(hf_state_dict) elif not skip_ngram_state and ple is not None: ple.ple_embedding.export_table_to_hf(hf_state_dict) + if not to_mcore and not skip_ngram_state and self.pp_size > 1: + # The shards are assembled on tp rank 0 of the stage owning the PLE + # layer -- which is exactly `pp_src_rank` within the tp0 pp group. + # Mirror the ngram buffers above so every pp rank (in particular the + # master that writes the checkpoint) carries them; groups whose tp + # coord does not own the table receive None and insert nothing. + shard_prefix = 'ple.ple_embedding.ngram_embedding' + keys = [f'{shard_prefix}.shard_{i}.weight' for i in range(self.config.split_ngram_parts)] + keys.append(Qwen4ExpTextNGramEmbedding._NGRAM_SCALE_KEY) + for key in keys: + obj = [hf_state_dict.get(key)] + dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) + if obj[0] is not None: + hf_state_dict[key] = obj[0] self._converting_ple = True try: for mg_key, hf_key in [('key_proj.weight', 'ple.key_proj.weight'), diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index 17bdbb14..e84efd14 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -17,6 +17,9 @@ from .kernels import gather_ple_rows, ple_gate_conv_triton _MASK64 = (1 << 64) - 1 +# Cap the per-all_reduce GPU buffer when exporting the (potentially huge) ngram +# table. +_EXPORT_CHUNK_BYTES = 256 << 20 _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 _SPLITMIX_M2 = 0x94D049BB133111EB @@ -219,16 +222,26 @@ def fill_table_from_hf(self, hf_state_dict): @torch.no_grad() def export_table_to_hf(self, hf_state_dict, prefix=''): - """Reverse of ``fill_table_from_hf``: write the offloaded table back as HF - shards so a full-parameter checkpoint is self-contained. + """Reverse of ``fill_table_from_hf``: write the n-gram table back as HF + shards so a full-parameter checkpoint is self-contained. Works for both + the host-resident (PLE_CPU_OFFLOAD=1) table and the TP-sharded + ``VocabParallelEmbedding`` (the default, which receives gradients and so + must not be silently dropped from the checkpoint). """ - if not self.cpu_offload: - return total = self.padded_vocab_size parts = self.split_ngram_parts shard_size = (total + parts - 1) // parts tp_rank = parallel_state.get_tensor_model_parallel_rank() tp_group = parallel_state.get_tensor_model_parallel_group() + if self.cpu_offload: + table, table_start, table_end, dtype = (self.host_table, self.vocab_start, self.vocab_end, + self.host_table.dtype) + else: + emb = self.ngram_embedding + per_partition = emb.num_embeddings_per_partition + table_start = tp_rank * per_partition + table_end = min(total, table_start + per_partition) + table, dtype = emb.weight.data, emb.weight.dtype # Inverse of fill_table_from_hf: the checkpoint format is fp8 shards + # scalar `weight_scale`, so divide by the scale stashed during loading and # cast back to fp8. Without a known scale the values cannot be represented @@ -237,24 +250,39 @@ def export_table_to_hf(self, hf_state_dict, prefix=''): if scale is None: get_logger().warning(f'`{self._NGRAM_SCALE_KEY}` was not seen during loading; exporting the PLE ngram ' 'embedding without re-quantizing to fp8.') + # Reduce on GPU: NCCL has no CPU backend. Each rank scatters its owned rows + # into a full shard chunk, sums across TP (rows are disjoint, so sum == + # gather), then rank 0 keeps the CPU copy. Chunking bounds the GPU + # transient: the shard can be far larger than the free VRAM left while the + # training state is still resident. + device = torch.cuda.current_device() + tp_size = parallel_state.get_tensor_model_parallel_world_size() + elem_size = torch.tensor([], dtype=dtype).element_size() + chunk_rows = max(1, _EXPORT_CHUNK_BYTES // max(1, table.shape[-1] * elem_size)) + out_dtype = torch.float8_e4m3fn if scale is not None else dtype for i in range(parts): cs, ce = i * shard_size, min((i + 1) * shard_size, total) - # Reduce on GPU: NCCL has no CPU backend, and the host table is pinned - # CPU. Each rank scatters its owned rows into a full shard, sums across - # TP (rows are disjoint, so sum == gather), then rank 0 keeps the CPU copy. - device = torch.cuda.current_device() - local = torch.zeros(ce - cs, self.host_table.shape[-1], dtype=self.host_table.dtype, device=device) - s, e = max(cs, self.vocab_start), min(ce, self.vocab_end) - if s < e: - local[s - cs:e - cs] = self.host_table[s - self.vocab_start:e - self.vocab_start].to(device) - if self._tp_size > 1: - torch.distributed.all_reduce(local, group=tp_group) + # Accumulate the shard on CPU chunk by chunk; assigning inside the + # loop would keep only the last chunk of each shard. + shard = None + if tp_rank == 0: + shard = torch.zeros(ce - cs, table.shape[-1], dtype=out_dtype, device='cpu') + for start in range(cs, ce, chunk_rows): + end = min(ce, start + chunk_rows) + local = torch.zeros(end - start, table.shape[-1], dtype=dtype, device=device) + s, e = max(start, table_start), min(end, table_end) + if s < e: + local[s - start:e - start] = table[s - table_start:e - table_start].to(device) + if tp_size > 1: + torch.distributed.all_reduce(local, group=tp_group) + if tp_rank == 0: + # Re-quantize after the all_reduce: fp8 is not a valid accumulation + # dtype for NCCL, and the sum must happen in the loaded dtype. + if scale is not None: + local = (local.to(torch.float32) / scale.to(local.device)).to(torch.float8_e4m3fn) + shard[start - cs:end - cs] = local.cpu() if tp_rank == 0: - # Re-quantize after the all_reduce: fp8 is not a valid accumulation - # dtype for NCCL, and the sum must happen in the loaded dtype. - if scale is not None: - local = (local.to(torch.float32) / scale.to(local.device)).to(torch.float8_e4m3fn) - hf_state_dict[f'{prefix}ple.ple_embedding.ngram_embedding.shard_{i}.weight'] = local.cpu() + hf_state_dict[f'{prefix}ple.ple_embedding.ngram_embedding.shard_{i}.weight'] = shard if tp_rank == 0 and scale is not None: key = f'{prefix}{self._NGRAM_SCALE_KEY}' if key not in hf_state_dict: From d3cfd8073f676dc3a7e5cf5d19dff2a6fdef3a82 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Mon, 14 Sep 2026 22:21:31 +0800 Subject: [PATCH 2/6] revert --- src/mcore_bridge/bridge/gpt_bridge.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index d2c2c43c..acfff06c 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -849,9 +849,7 @@ def _get_hf_experts_attr(self, is_mtp: bool = False): if (self._is_saving and not is_mtp and not self.config.fp8_param and not self._peft_format and self.model_type == 'qwen3_5_moe'): return True, True - if self.model_type in { - 'glm4v_moe', 'glm5_next', 'kimi_vl', 'qwen3_omni_moe', 'qwen3_5_moe' - } or self.llm_model_type in { + if self.model_type in {'glm4v_moe', 'kimi_vl', 'qwen3_omni_moe', 'qwen3_5_moe'} or self.llm_model_type in { 'qwen2_moe', 'qwen3_moe', 'deepseek_v2', 'deepseek_v3', 'kimi_k2', 'dots1', 'ernie4_5_moe', 'glm4_moe', 'glm4_moe_lite', 'minimax_m2', 'olmoe', 'qwen3_next', 'glm_moe_dsa', 'deepseek_v32', 'deepseek_v4' }: From 71c6073b31a4bbf6ce086190131a3fa69a7fd7b3 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Tue, 15 Sep 2026 12:46:05 +0800 Subject: [PATCH 3/6] fix(qwen4_exp): chunked export to bound GPU peaks; export GPU-resident PLE table; cross-PP PLE shard transfer --- src/mcore_bridge/bridge/gpt_bridge.py | 5 ++- src/mcore_bridge/model/gpts/qwen4_exp.py | 48 ++++++++++++++++++++---- src/mcore_bridge/model/modules/ple.py | 3 +- src/mcore_bridge/utils/constants.py | 6 +++ 4 files changed, 52 insertions(+), 10 deletions(-) create mode 100644 src/mcore_bridge/utils/constants.py diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index acfff06c..145729d9 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -18,6 +18,7 @@ from mcore_bridge.tuners import LoraParallelLinear from mcore_bridge.utils import (MxFp4Dequantizer, PackedDequantizer, SafetensorLazyLoader, StreamingSafetensorSaver, deep_getattr, gc_collect, get_logger, is_master, unwrap_model) +from mcore_bridge.utils.constants import EXPORT_CHUNK_BYTES logger = get_logger() @@ -30,7 +31,7 @@ class GPTBridge: fp8_block_size = 128 # Bound the per-collective GPU buffer when a full gathered tensor is # streamed to CPU (checkpoint save / CPU-offloaded weight sync). - export_chunk_bytes = 256 << 20 + export_chunk_bytes = EXPORT_CHUNK_BYTES hf_layers_prefix = 'model.layers' hf_mtp_prefix = 'model.layers' hf_embed_key = 'model.embed_tokens.weight' @@ -511,6 +512,8 @@ def _get_weight( if tensor.dtype == torch.uint8: mg_scale_inv = self._all_gather_tp(mg_scale_inv, tp_dim, is_expert) mg_scale_inv = self._broadcast_ep_pp(mg_scale_inv, is_expert) + if mg_scale_inv is not None and mg_scale_inv.device != tensor.device: + mg_scale_inv = mg_scale_inv.to(tensor.device) tensor = tensor.view(torch.float8_e4m3fn) assert tensor is not None, f'mg_key: {mg_key}' if offset: diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index e8c08582..63bfebca 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -351,6 +351,33 @@ def _get_pp_src_rank(self, has_module: bool) -> int: dist.all_reduce(holder, op=dist.ReduceOp.MAX, group=self.pp_group) return int(holder.item()) + def _broadcast_pp_weight(self, tensor, pp_src_rank: int): + """Cross-pp transfer of one exported weight through the tp-aligned pp + group. `tensor` is non-None only on the exporting rank (tp rank 0 of the + stage owning the PLE layer, i.e. the src of its pp group); the other pp + members receive it. The payload is streamed through + `_chunked_broadcast_pp` so no full-size GPU buffer is materialized, and + fp8 rides as uint8 (NCCL has no float8 dtype). Groups whose src has + nothing to transfer (tp != 0 coords) exchange only the empty meta. + """ + meta = [None if tensor is None else [list(tensor.shape), str(tensor.dtype).replace('torch.', '')]] + dist.broadcast_object_list(meta, src=pp_src_rank, group=self.pp_group) + if meta[0] is None: + return None + shape, dtype_name = meta[0] + dtype = getattr(torch, dtype_name) + as_uint8 = dtype == torch.float8_e4m3fn + if tensor is not None: + src_tensor = tensor if tensor.is_contiguous() else tensor.contiguous() + if as_uint8: + src_tensor = src_tensor.view(torch.uint8) + self._chunked_broadcast_pp(src_tensor, list(src_tensor.shape), src_tensor.dtype, pp_src_rank, + self.pp_group) + return tensor + recv_shape, recv_dtype = (shape, torch.uint8) if as_uint8 else (shape, dtype) + out = self._chunked_broadcast_pp(None, recv_shape, recv_dtype, pp_src_rank, self.pp_group) + return out.view(dtype) if as_uint8 else out + def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): ple = None if mg_layer is None else getattr(mg_layer, 'ple', None) if to_mcore: @@ -395,17 +422,22 @@ def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): if not to_mcore and not skip_ngram_state and self.pp_size > 1: # The shards are assembled on tp rank 0 of the stage owning the PLE # layer -- which is exactly `pp_src_rank` within the tp0 pp group. - # Mirror the ngram buffers above so every pp rank (in particular the - # master that writes the checkpoint) carries them; groups whose tp - # coord does not own the table receive None and insert nothing. + # Stream each shard to the other pp members (in particular the + # master that writes the checkpoint); groups whose tp coord does + # not own the table exchange only the empty meta. shard_prefix = 'ple.ple_embedding.ngram_embedding' keys = [f'{shard_prefix}.shard_{i}.weight' for i in range(self.config.split_ngram_parts)] - keys.append(Qwen4ExpTextNGramEmbedding._NGRAM_SCALE_KEY) for key in keys: - obj = [hf_state_dict.get(key)] - dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) - if obj[0] is not None: - hf_state_dict[key] = obj[0] + tensor = hf_state_dict.get(key) + transferred = self._broadcast_pp_weight(tensor, pp_src_rank) + if tensor is None and transferred is not None: + hf_state_dict[key] = transferred + # The scale is a tiny scalar; a pickled broadcast is fine. + key = Qwen4ExpTextNGramEmbedding._NGRAM_SCALE_KEY + obj = [hf_state_dict.get(key)] + dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) + if obj[0] is not None: + hf_state_dict[key] = obj[0] self._converting_ple = True try: for mg_key, hf_key in [('key_proj.weight', 'ple.key_proj.weight'), diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index e84efd14..e37b0336 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -12,6 +12,7 @@ from typing import List, Optional from ...utils import get_env_args, get_logger +from ...utils.constants import EXPORT_CHUNK_BYTES from ...utils.megatron_utils import get_num_samples, reconstruct_tensor_cp, split_cp_inputs from .hyper_connection_gated import Qwen4ExpTextGroupedRMSNorm from .kernels import gather_ple_rows, ple_gate_conv_triton @@ -19,7 +20,7 @@ _MASK64 = (1 << 64) - 1 # Cap the per-all_reduce GPU buffer when exporting the (potentially huge) ngram # table. -_EXPORT_CHUNK_BYTES = 256 << 20 +_EXPORT_CHUNK_BYTES = EXPORT_CHUNK_BYTES _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 _SPLITMIX_M2 = 0x94D049BB133111EB diff --git a/src/mcore_bridge/utils/constants.py b/src/mcore_bridge/utils/constants.py new file mode 100644 index 00000000..ab4bdbf0 --- /dev/null +++ b/src/mcore_bridge/utils/constants.py @@ -0,0 +1,6 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared constants.""" + +# Bound the per-collective GPU buffer when a full gathered tensor is streamed +# to CPU (checkpoint save / CPU-offloaded weight sync). +EXPORT_CHUNK_BYTES = 256 << 20 From 76917fcbc41ba2d05511a51756571cad7b6039c6 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Tue, 15 Sep 2026 13:20:09 +0800 Subject: [PATCH 4/6] perf(export): stream PLE table shards one at a time to bound host peaks - ple.py: split export_table_to_hf into a generator (iter_export_table_to_hf) that yields one shard at a time; the bulk API remains as a thin wrapper - qwen4_exp.py: stage the shard iterator in _set_layer_ple and interleave the cross-PP broadcast per shard (_iter_ple_table_export); consumed at _convert's yield point so no rank accumulates the 100GB-scale table in host memory - gpt_bridge.py: drain the staged iterator at the per-layer yield point and make the per-layer yield lazy Verified: 13 unit tests (#192's test_ple_checkpoint included), medium e2e TP2xPP2xEP2xETP2 (0 missing keys), and a load->export roundtrip that is bit-exact on all 16 PLE shards and 12 sampled other weights --- src/mcore_bridge/bridge/gpt_bridge.py | 17 ++- src/mcore_bridge/model/gpts/qwen4_exp.py | 70 ++++++++---- src/mcore_bridge/model/modules/ple.py | 131 +++++++++++++---------- 3 files changed, 135 insertions(+), 83 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 145729d9..ee45c2f6 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -1857,6 +1857,7 @@ def _convert_hf_state_dict(self, hf_state_dict, to_mcore): return res def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqdm_desc: str = 'Converting: '): + self._pending_export_iter = None if to_mcore: hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) hf_state_dict = self._convert_hf_state_dict(hf_state_dict, to_mcore) @@ -1906,7 +1907,11 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd yield else: res = self._convert_hf_state_dict(res, to_mcore) - yield from list(self._add_prefix(res, hf_prefix).items()) + # Drain any staged PLE table shards first: they are produced one + # at a time (see _iter_ple_table_export) so the 100GB-scale table + # never accumulates in host memory. + yield from self._drain_pending_export(hf_prefix) + yield from self._add_prefix(res, hf_prefix).items() hf_state_dict = {} if (not to_mcore or is_pp_last_stage) and self.config.mtp_num_layers: @@ -1933,6 +1938,16 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd yield from list(self._add_prefix(hf_state_dict, hf_prefix).items()) prog_bar.close() + def _drain_pending_export(self, hf_prefix: str): + """Yield (and release) the PLE table shards staged by _set_layer_ple.""" + it, self._pending_export_iter = self._pending_export_iter, None + if it is None: + return + for k, v in it: + if v is None: + continue + yield from self._add_prefix(self._convert_hf_state_dict({k: v}, False), hf_prefix).items() + def _convert_mtp_extra(self, mtp_layer, hf_state_dict, to_mcore, origin_hf_state_dict): for key in ['enorm.weight', 'hnorm.weight', 'eh_proj.weight']: self._set_state_dict(mtp_layer, key, hf_state_dict, key, to_mcore) diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 63bfebca..199a492f 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -378,7 +378,43 @@ def _broadcast_pp_weight(self, tensor, pp_src_rank: int): out = self._chunked_broadcast_pp(None, recv_shape, recv_dtype, pp_src_rank, self.pp_group) return out.view(dtype) if as_uint8 else out - def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): + def _iter_ple_table_export(self, ple, pp_src_rank, layer_prefix: str): + """Lazily export the PLE table shards one by one. On the exporting rank + (tp rank 0 of the owning stage) each shard is assembled by the chunked + all_reduce inside ``iter_export_table_to_hf`` and immediately streamed + to the other pp members; on the other stages it is received as it + arrives. Nothing accumulates: only the consumer (the safetensors + writer) drives the pace, so no rank ever holds more than a single shard + of the 100GB-scale table in host memory. The per-shard collectives keep + every rank of the tp-aligned pp groups in lockstep, exactly like the + synchronous version this generator replaces. + """ + parts = self.config.split_ngram_parts + scale_key = Qwen4ExpTextNGramEmbedding._NGRAM_SCALE_KEY + shard_prefix = 'ple.ple_embedding.ngram_embedding' + # On tp != 0 ranks of the owning stage the table iterator executes the + # same all_reduces but yields nothing (shards exist only on tp rank 0). + table_iter = ple.ple_embedding.iter_export_table_to_hf() if ple is not None else iter(()) + for i in range(parts): + shard = next(table_iter, (None, None))[1] if ple is not None else None + if self.pp_size > 1: + shard = self._broadcast_pp_weight(shard, pp_src_rank) + if shard is not None: + yield f'{layer_prefix}{shard_prefix}.shard_{i}.weight', shard + # The scale is a tiny scalar; a pickled broadcast is fine. + scale = None + if ple is not None: + for k, v in table_iter: + if k == scale_key: + scale = v + if self.pp_size > 1: + obj = [scale] + dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) + scale = obj[0] + if scale is not None: + yield f'{layer_prefix}{scale_key}', scale + + def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool, layer_prefix: str = ''): ple = None if mg_layer is None else getattr(mg_layer, 'ple', None) if to_mcore: # Only the stage owning the PLE layer reaches this path, so it @@ -417,27 +453,15 @@ def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): if not skip_ngram_state and to_mcore: # The table's only ingestion path: fill from the HF checkpoint shards. ple.ple_embedding.fill_table_from_hf(hf_state_dict) - elif not skip_ngram_state and ple is not None: - ple.ple_embedding.export_table_to_hf(hf_state_dict) - if not to_mcore and not skip_ngram_state and self.pp_size > 1: - # The shards are assembled on tp rank 0 of the stage owning the PLE - # layer -- which is exactly `pp_src_rank` within the tp0 pp group. - # Stream each shard to the other pp members (in particular the - # master that writes the checkpoint); groups whose tp coord does - # not own the table exchange only the empty meta. - shard_prefix = 'ple.ple_embedding.ngram_embedding' - keys = [f'{shard_prefix}.shard_{i}.weight' for i in range(self.config.split_ngram_parts)] - for key in keys: - tensor = hf_state_dict.get(key) - transferred = self._broadcast_pp_weight(tensor, pp_src_rank) - if tensor is None and transferred is not None: - hf_state_dict[key] = transferred - # The scale is a tiny scalar; a pickled broadcast is fine. - key = Qwen4ExpTextNGramEmbedding._NGRAM_SCALE_KEY - obj = [hf_state_dict.get(key)] - dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) - if obj[0] is not None: - hf_state_dict[key] = obj[0] + if not to_mcore and not skip_ngram_state: + # Stream the table shards one at a time instead of accumulating the + # 100GB-scale table in host memory: the generator below is only + # consumed at _convert's yield point (the safetensors writer drives + # the pace), and each shard is broadcast to the other pp members as + # it is produced, then released once written. + self._pending_export_iter = self._iter_ple_table_export(ple, pp_src_rank, layer_prefix) + else: + self._pending_export_iter = None self._converting_ple = True try: for mg_key, hf_key in [('key_proj.weight', 'ple.key_proj.weight'), @@ -460,7 +484,7 @@ def _set_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: i hf_state_dict.update(self._set_layer_mlp(mg_layer, hf_state_dict, layer_idx, to_mcore)) self._set_layer_hc(mg_layer, hf_state_dict, to_mcore) if (layer_idx + 1) in (self.config.ple_layer_ids or []): - self._set_layer_ple(mg_layer, hf_state_dict, to_mcore) + self._set_layer_ple(mg_layer, hf_state_dict, to_mcore, layer_prefix=hf_prefix) if to_mcore: hf_state_dict = {} else: diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index 483519ff..b1f67db2 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -236,72 +236,85 @@ def fill_table_from_hf(self, hf_state_dict): else: emb.weight.data[s - tp_start:e - tp_start] = weight[s - cs:e - cs].to(dtype) + @torch.no_grad() + def iter_export_table_to_hf(self, prefix=''): + """Generator variant of ``export_table_to_hf``: produce one shard at a + time so the 100GB-scale table never accumulates on the exporting rank -- + the consumer (the safetensors writer) drives the pace and each shard can + be released once written. Non-zero TP ranks execute the same all_reduces + but yield nothing. The collectives advance only while the generator is + being consumed. + """ + with torch.no_grad(): + total = self.padded_vocab_size + parts = self.split_ngram_parts + shard_size = (total + parts - 1) // parts + tp_rank = parallel_state.get_tensor_model_parallel_rank() + tp_group = parallel_state.get_tensor_model_parallel_group() + tp_size = parallel_state.get_tensor_model_parallel_world_size() + if self.cpu_offload: + table = self.host_table + tp_start, tp_end = self.vocab_start, self.vocab_end + device = torch.cuda.current_device() + else: + table = self.ngram_embedding.weight + tp_start = tp_rank * self.ngram_embedding.num_embeddings_per_partition + tp_end = min(tp_start + table.shape[0], total) + device = table.device + # Inverse of fill_table_from_hf: the checkpoint format is fp8 shards + + # scalar `weight_scale`, so divide by the scale stashed during loading + # and cast back to fp8. Without a known scale the values cannot be + # represented as fp8 + scale; keep the current dtype and warn. + scale = getattr(self, '_ngram_weight_scale', None) + if scale is None: + get_logger().warning(f'`{self._NGRAM_SCALE_KEY}` was not seen during loading; exporting the PLE ngram ' + 'embedding without re-quantizing to fp8.') + # Reduce on GPU: NCCL has no CPU backend, and the host table is pinned + # CPU. Each rank scatters its owned rows into shard chunks, sums across + # TP (rows are disjoint, so sum == gather), then rank 0 assembles the + # shard on CPU chunk by chunk. Chunking bounds the GPU transient: a + # shard can be far larger than the free VRAM left while the training + # state is still resident. + elem_size = torch.tensor([], dtype=table.dtype).element_size() + chunk_rows = max(1, _EXPORT_CHUNK_BYTES // max(1, table.shape[-1] * elem_size)) + out_dtype = torch.float8_e4m3fn if scale is not None else table.dtype + for i in range(parts): + cs, ce = i * shard_size, min((i + 1) * shard_size, total) + # Accumulate the shard on CPU chunk by chunk; assigning inside the + # loop would keep only the last chunk of each shard. + shard = None + if tp_rank == 0: + shard = torch.zeros(ce - cs, table.shape[-1], dtype=out_dtype, device='cpu') + for start in range(cs, ce, chunk_rows): + end = min(ce, start + chunk_rows) + local = torch.zeros(end - start, table.shape[-1], dtype=table.dtype, device=device) + s, e = max(start, tp_start), min(end, tp_end) + if s < e: + local[s - start:e - start] = table[s - tp_start:e - tp_start].to(device) + if tp_size > 1: + torch.distributed.all_reduce(local, group=tp_group) + if tp_rank == 0: + # Re-quantize after the all_reduce: fp8 is not a valid accumulation + # dtype for NCCL, and the sum must happen in the loaded dtype. + if scale is not None: + local = (local.to(torch.float32) / scale.to(local.device)).to(torch.float8_e4m3fn) + shard[start - cs:end - cs] = local.cpu() + if tp_rank == 0: + yield f'{prefix}ple.ple_embedding.ngram_embedding.shard_{i}.weight', shard + if tp_rank == 0 and scale is not None: + yield f'{prefix}{self._NGRAM_SCALE_KEY}', scale.reshape(()) + @torch.no_grad() def export_table_to_hf(self, hf_state_dict, prefix=''): """Reverse of ``fill_table_from_hf``: write the current table back as HF shards so a full-parameter checkpoint is self-contained. Works for both the host-resident (``PLE_CPU_OFFLOAD=1``) table and the TP-sharded ``VocabParallelEmbedding`` (the default, which receives gradients and so - must not be silently dropped from the checkpoint). + must not be silently dropped from the checkpoint). Bulk API: consumes + :meth:`iter_export_table_to_hf` into ``hf_state_dict``. """ - total = self.padded_vocab_size - parts = self.split_ngram_parts - shard_size = (total + parts - 1) // parts - tp_rank = parallel_state.get_tensor_model_parallel_rank() - tp_group = parallel_state.get_tensor_model_parallel_group() - tp_size = parallel_state.get_tensor_model_parallel_world_size() - if self.cpu_offload: - table = self.host_table - tp_start, tp_end = self.vocab_start, self.vocab_end - device = torch.cuda.current_device() - else: - table = self.ngram_embedding.weight - tp_start = tp_rank * self.ngram_embedding.num_embeddings_per_partition - tp_end = min(tp_start + table.shape[0], total) - device = table.device - # Inverse of fill_table_from_hf: the checkpoint format is fp8 shards + - # scalar `weight_scale`, so divide by the scale stashed during loading and - # cast back to fp8. Without a known scale the values cannot be represented - # as fp8 + scale; keep the current dtype and warn. - scale = getattr(self, '_ngram_weight_scale', None) - if scale is None: - get_logger().warning(f'`{self._NGRAM_SCALE_KEY}` was not seen during loading; exporting the PLE ngram ' - 'embedding without re-quantizing to fp8.') - # Reduce on GPU: NCCL has no CPU backend, and the host table is pinned - # CPU. Each rank scatters its owned rows into shard chunks, sums across TP - # (rows are disjoint, so sum == gather), then rank 0 assembles the shard on - # CPU chunk by chunk. Chunking bounds the GPU transient: a shard can be far - # larger than the free VRAM left while the training state is still resident. - elem_size = torch.tensor([], dtype=table.dtype).element_size() - chunk_rows = max(1, _EXPORT_CHUNK_BYTES // max(1, table.shape[-1] * elem_size)) - out_dtype = torch.float8_e4m3fn if scale is not None else table.dtype - for i in range(parts): - cs, ce = i * shard_size, min((i + 1) * shard_size, total) - # Accumulate the shard on CPU chunk by chunk; assigning inside the - # loop would keep only the last chunk of each shard. - shard = None - if tp_rank == 0: - shard = torch.zeros(ce - cs, table.shape[-1], dtype=out_dtype, device='cpu') - for start in range(cs, ce, chunk_rows): - end = min(ce, start + chunk_rows) - local = torch.zeros(end - start, table.shape[-1], dtype=table.dtype, device=device) - s, e = max(start, tp_start), min(end, tp_end) - if s < e: - local[s - start:e - start] = table[s - tp_start:e - tp_start].to(device) - if tp_size > 1: - torch.distributed.all_reduce(local, group=tp_group) - if tp_rank == 0: - # Re-quantize after the all_reduce: fp8 is not a valid accumulation - # dtype for NCCL, and the sum must happen in the loaded dtype. - if scale is not None: - local = (local.to(torch.float32) / scale.to(local.device)).to(torch.float8_e4m3fn) - shard[start - cs:end - cs] = local.cpu() - if tp_rank == 0: - hf_state_dict[f'{prefix}ple.ple_embedding.ngram_embedding.shard_{i}.weight'] = shard - if tp_rank == 0 and scale is not None: - key = f'{prefix}{self._NGRAM_SCALE_KEY}' - if key not in hf_state_dict: - hf_state_dict[key] = scale.reshape(()) + for key, tensor in self.iter_export_table_to_hf(prefix): + hf_state_dict[key] = tensor def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor: # Mirrors transformers `_shift_right_ignore_eos`: segment-aware shift, From 3d0f1de8f7bc4f5cf545e5f0fb87a78623579cc4 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Tue, 15 Sep 2026 15:39:48 +0800 Subject: [PATCH 5/6] refactor(qwen4_exp): transport cross-pp PLE shards as raw bytes instead of hardcoding fp8 Any low-width dtype (fp8 e4m3/e5m2, int8, packed fp4, ...) now rides as a flattened uint8 payload; the meta carries the original shape/dtype and the receiver views back. Removes the float8_e4m3fn-only special case. --- src/mcore_bridge/model/gpts/qwen4_exp.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 199a492f..aa0c4dce 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import copy +import math import torch import torch.distributed as dist import torch.nn.functional as F @@ -357,8 +358,11 @@ def _broadcast_pp_weight(self, tensor, pp_src_rank: int): stage owning the PLE layer, i.e. the src of its pp group); the other pp members receive it. The payload is streamed through `_chunked_broadcast_pp` so no full-size GPU buffer is materialized, and - fp8 rides as uint8 (NCCL has no float8 dtype). Groups whose src has - nothing to transfer (tp != 0 coords) exchange only the empty meta. + it always rides as raw bytes (flattened uint8): any current or future + low-width dtype (fp8 e4m3/e5m2, int8, packed fp4, ...) is transported + without NCCL dtype concerns; the meta carries the original shape/dtype + for the receiver to view back. Groups whose src has nothing to transfer + (tp != 0 coords) exchange only the empty meta. """ meta = [None if tensor is None else [list(tensor.shape), str(tensor.dtype).replace('torch.', '')]] dist.broadcast_object_list(meta, src=pp_src_rank, group=self.pp_group) @@ -366,17 +370,13 @@ def _broadcast_pp_weight(self, tensor, pp_src_rank: int): return None shape, dtype_name = meta[0] dtype = getattr(torch, dtype_name) - as_uint8 = dtype == torch.float8_e4m3fn if tensor is not None: - src_tensor = tensor if tensor.is_contiguous() else tensor.contiguous() - if as_uint8: - src_tensor = src_tensor.view(torch.uint8) - self._chunked_broadcast_pp(src_tensor, list(src_tensor.shape), src_tensor.dtype, pp_src_rank, - self.pp_group) + payload = tensor.contiguous().flatten().view(torch.uint8) + self._chunked_broadcast_pp(payload, list(payload.shape), payload.dtype, pp_src_rank, self.pp_group) return tensor - recv_shape, recv_dtype = (shape, torch.uint8) if as_uint8 else (shape, dtype) - out = self._chunked_broadcast_pp(None, recv_shape, recv_dtype, pp_src_rank, self.pp_group) - return out.view(dtype) if as_uint8 else out + byte_count = math.prod(shape) * torch.empty((), dtype=dtype).element_size() + out = self._chunked_broadcast_pp(None, [byte_count], torch.uint8, pp_src_rank, self.pp_group) + return out.view(dtype).view(shape) def _iter_ple_table_export(self, ple, pp_src_rank, layer_prefix: str): """Lazily export the PLE table shards one by one. On the exporting rank From 71668a214ec90fa0f4cd575dd349bba332d219f4 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Tue, 15 Sep 2026 15:44:17 +0800 Subject: [PATCH 6/6] fix(export): guard _chunk_rows against empty dim0 and fix E125 lint - _chunk_rows returns max(1, min(rows, shape[0])) so a zero-row shape can never produce range(0, 0, 0) in _chunked_broadcast_pp (currently unreachable -- PLE shards have >= 1 row and _broadcast_pp_weight only calls in with a non-empty meta -- but a robustness gap, not a live bug) - extract elem_size in _broadcast_ep_pp's gate to fix flake8 E125 --- src/mcore_bridge/bridge/gpt_bridge.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index ee45c2f6..082c4c0b 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -357,7 +357,7 @@ def _chunk_rows(self, shape, elem_size: int) -> int: for s in shape[1:]: inner *= s rows = max(1, self.export_chunk_bytes // max(1, inner * elem_size)) - return min(rows, shape[0]) + return max(1, min(rows, shape[0])) def _chunked_all_gather_tp(self, tensor, tp_dim: int, tp_group, tp_size: int): """All-gather `tensor` along tp_dim and assemble the result in host memory @@ -457,8 +457,8 @@ def _broadcast_ep_pp(self, tensor, is_expert): shape = meta_data[1:1 + meta_data[0]].tolist() dtype = dtype_mapping[meta_data[-1].item()] numel = math.prod(shape) if shape else 1 - if self._stream_to_cpu() and numel * torch.empty( - (), dtype=dtype).element_size() > (self.export_chunk_bytes) and len(shape) > 0: + elem_size = torch.empty((), dtype=dtype).element_size() + if self._stream_to_cpu() and numel * elem_size > self.export_chunk_bytes and len(shape) > 0: return self._chunked_broadcast_pp(None, shape, dtype, src_rank, pp_group) tensor = torch.empty(shape, device='cuda', dtype=dtype) dist.broadcast(tensor, src=src_rank, group=pp_group)