diff --git a/ci/scripts/CI_ENV.sh b/ci/scripts/CI_ENV.sh index 7b4431e94..f6da84956 100644 --- a/ci/scripts/CI_ENV.sh +++ b/ci/scripts/CI_ENV.sh @@ -28,6 +28,7 @@ export ROLLOUT_DAPO_DATA_PATH=${CI_SHARE_DATA}/rl_test_judger_dapo_math_data.jso export GEO_ROLLOUT_DATA_PATH=${CI_SHARE_DATA}/rl_test_judge_geo_data.jsonl export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0 export XTUNER_DETERMINISTIC=true +export TORCHINDUCTOR_DYNAMIC_SCALE_RBLOCK=0 export XTUNER_USE_LMDEPLOY=1 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True export PYTHONPYCACHEPREFIX=/tmp diff --git a/xtuner/v1/data_proto/sequence_context.py b/xtuner/v1/data_proto/sequence_context.py index 6860ee362..fa1829a91 100644 --- a/xtuner/v1/data_proto/sequence_context.py +++ b/xtuner/v1/data_proto/sequence_context.py @@ -1,5 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. -import itertools from typing import cast import torch @@ -9,9 +8,6 @@ from .utils import gather_for_sequence_parallel, pad_to_multiple_of, split_for_sequence_parallel -_DSA_TOPK_CONTEXT_IDS = itertools.count() - - class DSATopKCacheState: """Mutable DSA cross-layer top-k cache, scoped to one microbatch. @@ -29,7 +25,7 @@ class DSATopKCacheState: offloaded: dict[int, str] # OffloadManager key for each CPU-resident source. released_sources: set[int] # Sources whose backward replay lifetime has ended. checkpoint_active: bool # Whether checkpoint forward retained this cache for replay. - context_id: int # Process-local identifier used to make offload keys unique. + offload_slot: int # Stable offload slot among concurrently active microbatches. mtp_forward_uses_remaining: dict[int, int] # Original-forward MTP uses left per shared source. mtp_replays_remaining: dict[int, int] # Backward MTP replays left per shared source. @@ -40,7 +36,7 @@ def __init__( offloaded: dict[int, str] | None = None, released_sources: set[int] | None = None, checkpoint_active: bool = False, - context_id: int | None = None, + offload_slot: int = 0, mtp_forward_uses_remaining: dict[int, int] | None = None, mtp_replays_remaining: dict[int, int] | None = None, ) -> None: @@ -50,7 +46,7 @@ def __init__( self.offloaded = {} if offloaded is None else offloaded self.released_sources = set() if released_sources is None else released_sources self.checkpoint_active = checkpoint_active - self.context_id = next(_DSA_TOPK_CONTEXT_IDS) if context_id is None else context_id + self.offload_slot = offload_slot self.mtp_forward_uses_remaining = {} if mtp_forward_uses_remaining is None else mtp_forward_uses_remaining self.mtp_replays_remaining = {} if mtp_replays_remaining is None else mtp_replays_remaining diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 23b2369ed..f27e0a2db 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -475,6 +475,14 @@ def post_micro_batch_forward(self, batch_outputs: Sequence[MoEModelOutputs]) -> moe_info = cast(MoEBatchForwardInfo, base_info) return moe_info + @staticmethod + def _prepare_seq_ctx_topk_cache(seq_ctx_list: Sequence[SequenceContext]) -> None: + # A slot owns both runtime offload state and pinned storage. TrainEngine + # finishes backward before the next accumulation group, so only contexts + # in one model call can be live concurrently and need distinct slots. + for offload_slot, seq_ctx in enumerate(seq_ctx_list): + seq_ctx.dsa_topk_cache.offload_slot = offload_slot + def _micro_batch_forward( self, seq_ctx_list: list[SequenceContext], @@ -486,6 +494,7 @@ def _micro_batch_forward( This method processes multiple micro-batches in parallel, similar to how MoEDecoderLayer handles micro-batching at the layer level. """ + self._prepare_seq_ctx_topk_cache(seq_ctx_list) if self.config.return_hidden_states: raise NotImplementedError @@ -624,7 +633,9 @@ def _micro_batch_forward( input_ids=seq_ctx.input_ids.clone() if seq_ctx.input_ids is not None else None, position_ids=seq_ctx.position_ids.clone(), inputs_embeds=seq_ctx.inputs_embeds.clone() if seq_ctx.inputs_embeds is not None else None, - dsa_topk_cache=DSATopKCacheState(), + dsa_topk_cache=DSATopKCacheState( + offload_slot=seq_ctx.dsa_topk_cache.offload_slot, + ), ) ) @@ -747,6 +758,7 @@ def _forward( loss_ctx: MoELossContextDict | None, return_router_logits: bool = False, ) -> MoEModelOutputs: + self._prepare_seq_ctx_topk_cache([seq_ctx]) input_ids = seq_ctx.input_ids position_ids = seq_ctx.position_ids @@ -854,7 +866,9 @@ def _forward( input_ids=input_ids.clone() if input_ids is not None else None, position_ids=position_ids.clone(), inputs_embeds=seq_ctx.inputs_embeds.clone() if seq_ctx.inputs_embeds is not None else None, - dsa_topk_cache=DSATopKCacheState(), + dsa_topk_cache=DSATopKCacheState( + offload_slot=seq_ctx.dsa_topk_cache.offload_slot, + ), ) # MTP uses its own mask; main mask's non-pad indices do not apply. mtp_nonpad_indices = torch.nonzero(mtp_seq_ctx.mask, as_tuple=True)[1] diff --git a/xtuner/v1/module/attention/dsa_topk_sharing.py b/xtuner/v1/module/attention/dsa_topk_sharing.py index ce6e2f574..7e0ca4fd2 100644 --- a/xtuner/v1/module/attention/dsa_topk_sharing.py +++ b/xtuner/v1/module/attention/dsa_topk_sharing.py @@ -106,7 +106,7 @@ def after_recompute_release(self, seq_ctx: SequenceContext, source_layer_idx: in seq_ctx.dsa_topk_cache.indices.pop(source_layer_idx, None) def _offload_key(self, seq_ctx: SequenceContext, source_layer_idx: int) -> str: - return f"dsa_topk_{seq_ctx.dsa_topk_cache.context_id}_{source_layer_idx}" + return f"dsa_topk_{seq_ctx.dsa_topk_cache.offload_slot}_{source_layer_idx}" class ActivationOffloadedTopKResidency(GpuTopKResidency): @@ -178,7 +178,16 @@ def after_original_forward_last_use(self, seq_ctx: SequenceContext, source_layer return key = self._offload_key(seq_ctx, source_layer_idx) - cpu_buffer = OffloadManager().get_or_create_pin_memory(key, topk_indices.shape, topk_indices.dtype) + # A slot/source pair owns both the runtime entry and its reusable pinned + # storage. Reusing it while still live would overwrite the earlier D2H + # result, so fail loudly if the scheduling contract is violated. + if OffloadManager().has_runtime_key(key): + raise RuntimeError(f"DSA top-k offload slot is still active: {key}") + cpu_buffer = OffloadManager().get_or_create_pin_memory( + key, + topk_indices.shape, + topk_indices.dtype, + ) swap_tensor = SwapTensor(topk_indices, key, tensor_cpu=cpu_buffer) stream = self._stream_for_device(topk_indices.device) stream.wait_stream(torch.cuda.current_stream(topk_indices.device)) diff --git a/xtuner/v1/ops/sparse_mla/pytorch.py b/xtuner/v1/ops/sparse_mla/pytorch.py index 7813973cc..0a0392bde 100644 --- a/xtuner/v1/ops/sparse_mla/pytorch.py +++ b/xtuner/v1/ops/sparse_mla/pytorch.py @@ -68,7 +68,7 @@ def torch_dsa_topk_indices( topk = min(index_topk, kv_len) topk_scores, topk_indices = index_scores.topk(topk, dim=-1) topk_indices = topk_indices.masked_fill(topk_scores == -torch.inf, -1) - return topk_indices.squeeze(0).unsqueeze(1) + return topk_indices.squeeze(0).unsqueeze(1).to(torch.int32) def _packed_causal_mask(seq_ctx: SequenceContext, query_len: int, kv_len: int, device: torch.device) -> torch.Tensor: diff --git a/xtuner/v1/ops/sparse_mla/tilelang.py b/xtuner/v1/ops/sparse_mla/tilelang.py index 603e75a4f..d186942fe 100644 --- a/xtuner/v1/ops/sparse_mla/tilelang.py +++ b/xtuner/v1/ops/sparse_mla/tilelang.py @@ -170,7 +170,7 @@ def _tilelang_dsa_topk_indices_from_ranges( topk = min(index_topk, k.shape[0]) topk_scores, topk_indices = logits.topk(topk, dim=-1) topk_indices = topk_indices.masked_fill(topk_scores == -torch.inf, -1) - return topk_indices.to(torch.int64).unsqueeze(1) + return topk_indices.to(torch.int32).unsqueeze(1) @_tilelang_dsa_topk_indices_from_ranges.register_fake @@ -183,7 +183,7 @@ def _( index_topk: int, ) -> Tensor: topk = min(index_topk, k.shape[0]) - return torch.empty((q.shape[0], 1, topk), device=q.device, dtype=torch.int64) + return torch.empty((q.shape[0], 1, topk), device=q.device, dtype=torch.int32) @functools.cache diff --git a/xtuner/v1/utils/activation_offload.py b/xtuner/v1/utils/activation_offload.py index eac495b12..e0788af6e 100644 --- a/xtuner/v1/utils/activation_offload.py +++ b/xtuner/v1/utils/activation_offload.py @@ -229,6 +229,10 @@ def assert_exist(self, key): def exist(self, key): return key in self.items + def has_runtime_key(self, key): + """Return whether an offload key still owns a live runtime entry.""" + return key in self.items or key in self.may_npu_tensors + def assert_not_exist(self, key): if key not in self.items: raise RuntimeError(f"Key {key} already exist in items") diff --git a/xtuner/v1/utils/misc.py b/xtuner/v1/utils/misc.py index aad294a4a..4d3aede81 100644 --- a/xtuner/v1/utils/misc.py +++ b/xtuner/v1/utils/misc.py @@ -29,17 +29,34 @@ def set_deterministic(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" - # Inductor 会在 torch.compile 前读取 dynamic_scale_rblock;确定性模式必须尽早关闭。 - # torch.use_deterministic_algorithms(True) 只会让 reduction 的初始候选收敛成一个 config; - # dynamic rblock 仍可能在 precompile 后追加一个缩小 R*_BLOCK 的 launcher,并由 runtime - # benchmark 在多个 launcher 中选择。不同 rank/run 一旦选到不同 reduction 分块,浮点累加 - # 顺序就会变化,最终可能得到 bitwise 不同的梯度。 + # Inductor 有两个独立的 reduction 调优入口:dynamic rblock 会在 precompile 后调整 + # launcher,而 INNER reduction 会在 deterministic 检查前返回 persistent/contiguous + # 多个候选,由 runtime benchmark 选出配置。确定性模式必须同时关闭前者并固定后者。 + # Triton compile subprocess 不会继承 Python monkey patch,因此也要在当前进程串行编译, + # 避免不同 rank/run 采用不同的浮点累加顺序。 os.environ["TORCHINDUCTOR_DYNAMIC_SCALE_RBLOCK"] = "0" + os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1" from torch._inductor import config as inductor_config + from torch._inductor.runtime import triton_heuristics inductor_config.dynamic_scale_rblock = False + inductor_config.compile_threads = 1 torch.use_deterministic_algorithms(True, warn_only=True) + original_reduction_configs = triton_heuristics._reduction_configs + if not getattr(original_reduction_configs, "_xtuner_deterministic_patched", False): + + def deterministic_reduction_configs(*args, **kwargs): + return original_reduction_configs(*args, **kwargs)[:1] + + setattr(deterministic_reduction_configs, "_xtuner_deterministic_patched", True) + setattr( + deterministic_reduction_configs, + "_xtuner_original_reduction_configs", + original_reduction_configs, + ) + triton_heuristics._reduction_configs = deterministic_reduction_configs + # https://github.com/python/cpython/issues/82300#issuecomment-2169035092 if sys.version_info >= (3, 13):