From f62772679dfebaa875f0bcf3bc1750169f37fd57 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Mon, 27 Jul 2026 12:08:17 +0000 Subject: [PATCH 1/7] [Test] Add GLM-5.2 parallel HF round-trip coverage --- tests/engine/test_glm52_moe_train_engine.py | 84 ++++++++++++++++++++- 1 file changed, 82 insertions(+), 2 deletions(-) diff --git a/tests/engine/test_glm52_moe_train_engine.py b/tests/engine/test_glm52_moe_train_engine.py index baf40d3b6..d68ee89b8 100644 --- a/tests/engine/test_glm52_moe_train_engine.py +++ b/tests/engine/test_glm52_moe_train_engine.py @@ -2,6 +2,8 @@ TestGlm52OptimizedEngine test_sp2_ep4_micro2_compile_offload_train_step: SP2、EP4、micro2、compile 与双 offload 可联合训练。 +TestGlm52ParallelHFCheckpoint + test_fsdp2_ep4_mtp_hf_round_trip_preserves_weights: FSDP2、EP4、MTP 权重可经 HF 无损往返。 TestGlm52PretrainedEngine test_ep8_loss_curve_matches_reference: 预训练权重的 EP8 优化轨迹匹配数值基线。 test_tilewise_fp8_loss_curve_matches_bf16: tilewise FP8 训练轨迹接近 BF16。 @@ -10,6 +12,7 @@ test_dcp_round_trip_preserves_model_and_optimizer: DCP 往返保留模型与优化器状态。 """ +import json import math import os import shutil @@ -20,6 +23,7 @@ import torch import torch.distributed as dist +from safetensors import safe_open from torch.distributed.tensor import DTensor from torch.optim.lr_scheduler import LambdaLR @@ -104,7 +108,7 @@ def _tiny_checkpoint_config(dispatcher: str | None, ep_size: int) -> Glm52MoECon ) -def _tiny_sp_mtp_config() -> Glm52MoEConfig: +def _tiny_ep4_mtp_config() -> Glm52MoEConfig: config = _tiny_checkpoint_config(dispatcher="all2all", ep_size=4) config.hidden_size = 128 config.intermediate_size = 128 @@ -189,7 +193,7 @@ class TestGlm52OptimizedEngine(DeterministicDDPTestCase): def test_sp2_ep4_micro2_compile_offload_train_step(self): # 验证生产优化组合经两次梯度累积后 loss、梯度与优化器状态均有效。 self.create_pg("cuda") - model_cfg = _tiny_sp_mtp_config() + model_cfg = _tiny_ep4_mtp_config() engine = TrainEngine( model_cfg=model_cfg, optim_cfg=AdamWConfig(lr=1e-3, foreach=False), @@ -242,6 +246,82 @@ def world_size(self) -> int: return 8 +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestGlm52ParallelHFCheckpoint(DeterministicDDPTestCase): + def test_fsdp2_ep4_mtp_hf_round_trip_preserves_weights(self): + # 验证 FSDP2、EP4 分片的主干与共享 MTP 权重经公共 HF API 往返后完全一致。 + self.create_pg("cuda") + temp_dir = tempfile.mkdtemp() if dist.get_rank() == 0 else None + syncdir = [temp_dir] + dist.broadcast_object_list(syncdir, src=0) + first_hf_dir = Path(syncdir[0]) / "first" + second_hf_dir = Path(syncdir[0]) / "second" + + try: + model_cfg = _tiny_ep4_mtp_config() + model_cfg.compile_cfg = False + engine = TrainEngine( + model_cfg=model_cfg, + optim_cfg=AdamWConfig(lr=1e-3, foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, cpu_offload=False, torch_compile=False), + ) + engine.init_model_weights() + engine.save_hf(str(first_hf_dir)) + del engine + torch.cuda.empty_cache() + + restored_cfg = _tiny_ep4_mtp_config() + restored_cfg.compile_cfg = False + restored = TrainEngine( + model_cfg=restored_cfg, + optim_cfg=AdamWConfig(lr=1e-3, foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, cpu_offload=False, torch_compile=False), + ) + restored.from_hf(first_hf_dir, strict=True) + restored.save_hf(str(second_hf_dir)) + + if dist.get_rank() == 0: + with open(first_hf_dir / "model.safetensors.index.json") as f: + first_index = json.load(f)["weight_map"] + with open(second_hf_dir / "model.safetensors.index.json") as f: + second_index = json.load(f)["weight_map"] + + mtp_prefix = f"model.layers.{model_cfg.num_hidden_layers}." + self.assertTrue(any(key.startswith(mtp_prefix) for key in first_index)) + self.assertEqual(first_index.keys(), second_index.keys()) + first_files = {} + second_files = {} + for key in first_index: + first_filename = first_index[key] + second_filename = second_index[key] + if first_filename not in first_files: + first_files[first_filename] = safe_open( + first_hf_dir / first_filename, + framework="pt", + ) + if second_filename not in second_files: + second_files[second_filename] = safe_open( + second_hf_dir / second_filename, + framework="pt", + ) + self.assertTrue( + torch.equal( + first_files[first_filename].get_tensor(key), + second_files[second_filename].get_tensor(key), + ), + f"HF round-trip tensor mismatch: {key}", + ) + finally: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(syncdir[0], ignore_errors=True) + torch.cuda.empty_cache() + + @property + def world_size(self) -> int: + return 8 + + @unittest.skipUnless( torch.cuda.device_count() >= 8 and GLM5_2_TINY_MOE_PATH.exists(), f"requires 8 CUDA devices and GLM-5.2 checkpoint at {GLM5_2_TINY_MOE_PATH}", From e8e67c1238fd3c2598e276d7d5cd3e9b0635bb79 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Mon, 27 Jul 2026 12:08:59 +0000 Subject: [PATCH 2/7] Support Muon in GLM-5.2 SFT config --- examples/v1/config/sft_glm5p2.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/examples/v1/config/sft_glm5p2.py b/examples/v1/config/sft_glm5p2.py index eef2faba5..cab158ec2 100644 --- a/examples/v1/config/sft_glm5p2.py +++ b/examples/v1/config/sft_glm5p2.py @@ -1,6 +1,6 @@ import os -from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig, MuonConfig from xtuner.v1.datasets import OpenaiTokenizeFunctionConfig from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig from xtuner.v1.float8.config import Float8Config, ScalingGranularity @@ -91,11 +91,18 @@ def _get_float8_config() -> Float8Config | None: num_workers=int(os.environ.get("DATALOADER_NUM_WORKERS", "4")), ) -optim_cfg = AdamWConfig( - lr=float(os.environ.get("LR", "1e-6")), - foreach=_get_bool_env("ADAMW_FOREACH", False), - swap_optimizer=_get_bool_env("SWAP_OPTIMIZER", False), -) +lr = float(os.environ.get("LR", "1e-6")) +optimizer = os.environ.get("OPTIMIZER", "adamw").lower() +if optimizer == "muon": + optim_cfg = MuonConfig(lr=lr) +elif optimizer == "adamw": + optim_cfg = AdamWConfig( + lr=lr, + foreach=_get_bool_env("ADAMW_FOREACH", False), + swap_optimizer=_get_bool_env("SWAP_OPTIMIZER", False), + ) +else: + raise ValueError(f"Unsupported OPTIMIZER={optimizer!r}. Use adamw or muon.") lr_cfg = LRConfig(lr_type=os.environ.get("LR_TYPE", "cosine"), warmup_ratio=float(os.environ.get("WARMUP_RATIO", "0"))) fsdp_cfg = FSDPConfig( cpu_offload=_get_bool_env("CPU_OFFLOAD", False), From dc97cb3c24fe4cb9d3cbf19f8083f79d7cd17208 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 29 Jul 2026 03:52:25 +0000 Subject: [PATCH 3/7] [Fix] Freeze GLM-5.2 DSA indexer parameters --- tests/engine/test_glm52_moe_train_engine.py | 14 ++++++++++++++ xtuner/v1/module/attention/dsa_mla.py | 3 +++ 2 files changed, 17 insertions(+) diff --git a/tests/engine/test_glm52_moe_train_engine.py b/tests/engine/test_glm52_moe_train_engine.py index d68ee89b8..a10240338 100644 --- a/tests/engine/test_glm52_moe_train_engine.py +++ b/tests/engine/test_glm52_moe_train_engine.py @@ -409,12 +409,23 @@ def test_dcp_round_trip_preserves_model_and_optimizer(self): optim_cfg=AdamWConfig(), fsdp_cfg=FSDPConfig(cpu_offload=False, ep_size=2), ) + torch.manual_seed(0) engine.init_model_weights() with torch.no_grad(): for module in engine.model.modules(): if isinstance(module, NoAuxRouter): bias = module.e_score_correction_bias bias.copy_(torch.arange(bias.numel(), device=bias.device, dtype=bias.dtype)) + + input_ids = torch.arange(2, 12).view(1, -1) % config.vocab_size + seq_ctx = SequenceContext.from_input_ids((input_ids[:, :-1],), device=DEVICE) + data = {"seq_ctx": seq_ctx, "shifted_labels": input_ids[:, 1:]} + loss_ctx = engine.model.build_loss_ctx_batch([data], sp_mesh=None)[0] + engine.train_step([ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]) + grad_norm = engine.clip_grad_norm() + self.assertTrue(math.isfinite(float(grad_norm))) + engine.step_optimizer(grad_norm) + engine.save_dcp(weights_dir=weights_dir) dist.barrier() @@ -423,6 +434,9 @@ def test_dcp_round_trip_preserves_model_and_optimizer(self): optim_cfg=AdamWConfig(), fsdp_cfg=FSDPConfig(cpu_offload=False, ep_size=2), ) + # Frozen parameters are restored from the same base checkpoint in + # Trainer; matching initialization models that behavior here. + torch.manual_seed(0) restored.init_model_weights() restored.load_dcp(weights_dir=weights_dir) diff --git a/xtuner/v1/module/attention/dsa_mla.py b/xtuner/v1/module/attention/dsa_mla.py index 23e0f9f68..063e740b7 100644 --- a/xtuner/v1/module/attention/dsa_mla.py +++ b/xtuner/v1/module/attention/dsa_mla.py @@ -83,6 +83,9 @@ def __init__( self.k_norm = LayerNorm(index_head_dim, eps=1e-6) # weights_proj.weight: [index_n_heads, hidden_size] self.weights_proj = build_linear(hidden_size, index_n_heads, bias=False) + # The indexer only produces integer DSA top-k IDs under no_grad, so its + # parameters must not be registered with the training optimizer. + self.requires_grad_(False) @torch.no_grad() def forward( From cf95ca30f31347389dc9ee4dc572b5a1be59960d Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 29 Jul 2026 07:11:51 +0000 Subject: [PATCH 4/7] [Fix] Release cached memory after checkpoint resume DCP resume materializes temporary model and optimizer state that can remain cached until the first training step, starving NCCL allocations. Collect Python garbage and empty the device cache after the complete checkpoint state has been restored. --- tests/train/test_trainer.py | 20 +++++++++++++++++++- xtuner/v1/train/trainer.py | 4 ++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/train/test_trainer.py b/tests/train/test_trainer.py index f89a79168..0cb8853a0 100644 --- a/tests/train/test_trainer.py +++ b/tests/train/test_trainer.py @@ -958,10 +958,25 @@ def test_resume_and_load_checkpoint_cfg(tmp_path: Path): load_scheduler=False, ) + load_artifacts = {} + + def load_dcp(**_): + class TemporaryState: + pass + + temporary_state = TemporaryState() + temporary_state.cycle = temporary_state + load_artifacts["temporary_state"] = weakref.ref(temporary_state) + + if DEVICE == "cuda": + temporary_tensor = torch.empty(16 * 1024 * 1024, dtype=torch.uint8, device=DEVICE) + del temporary_tensor + load_artifacts["reserved_memory"] = torch.cuda.memory_reserved() + # 2. operate with ( patch.object(Dataloader, "load_state_dict") as mock_data_load_state_dict, - patch.object(FakeEngine, "load_dcp") as mock_load_dcp, + patch.object(FakeEngine, "load_dcp", side_effect=load_dcp) as mock_load_dcp, patch.object(SequentialLR, "load_state_dict") as mock_lr_load_state_dict, ): trainer = Trainer( @@ -987,6 +1002,9 @@ def test_resume_and_load_checkpoint_cfg(tmp_path: Path): load_states=True, load_args=True, ) + assert load_artifacts["temporary_state"]() is None + if DEVICE == "cuda": + assert torch.cuda.memory_reserved() < load_artifacts["reserved_memory"] # assert trainer._load_checkpoint_cfg.load_dataset is False # assert trainer._load_checkpoint_cfg.load_scheduler is False trainer.fit() diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index 1d279156d..ba3700db0 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -744,6 +744,10 @@ def __init__( if self._load_checkpoint_cfg.checkpoint_path is not None: self._load_checkpoint() + # DCP load materializes temporary model and optimizer state. Release + # its cached device memory before the first resumed training step. + gc.collect() + DEVICE_MODULE.empty_cache() self.hooks_config = self._setup_hooks(hooks_config=hooks_config) From 10df87928d0764e8fa17db2f50039bdda5fd3326 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 29 Jul 2026 07:20:14 +0000 Subject: [PATCH 5/7] [Enhance] Log memory after checkpoint resume cleanup Report per-rank allocated, reserved, free, and total device memory after checkpoint garbage collection and cache release so large-scale resume runs can verify the recovered headroom. --- tests/train/test_trainer.py | 3 ++- xtuner/v1/train/trainer.py | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/train/test_trainer.py b/tests/train/test_trainer.py index 0cb8853a0..a47ced1a4 100644 --- a/tests/train/test_trainer.py +++ b/tests/train/test_trainer.py @@ -900,7 +900,7 @@ def __call__(self, checkpoint, step, epoch, total_step, total_epoch): @patch("xtuner.v1.train.trainer.Trainer._prepare_model_input", Mock(return_value=[])) @patch("xtuner.v1.train.trainer.Trainer.build_engine", Mock(side_effect=lambda *args, **kwargs: FakeEngine())) -def test_resume_and_load_checkpoint_cfg(tmp_path: Path): +def test_resume_and_load_checkpoint_cfg(tmp_path: Path, capfd): # 0. prepare environment os.environ["LOCAL_RANK"] = "0" os.environ["RANK"] = "0" @@ -1005,6 +1005,7 @@ class TemporaryState: assert load_artifacts["temporary_state"]() is None if DEVICE == "cuda": assert torch.cuda.memory_reserved() < load_artifacts["reserved_memory"] + assert "[Checkpoint Resume Memory]" in capfd.readouterr().err # assert trainer._load_checkpoint_cfg.load_dataset is False # assert trainer._load_checkpoint_cfg.load_scheduler is False trainer.fit() diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index ba3700db0..f3e2b2f24 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -748,6 +748,14 @@ def __init__( # its cached device memory before the first resumed training step. gc.collect() DEVICE_MODULE.empty_cache() + if DEVICE != "cpu": + free_memory, total_memory = DEVICE_MODULE.mem_get_info() # type: ignore[attr-defined] + self.logger.info( + "[Checkpoint Resume Memory] " + f"allocated: {DEVICE_MODULE.memory_allocated() / (1024**3):.2f} GB, " # type: ignore[attr-defined] + f"reserved: {DEVICE_MODULE.memory_reserved() / (1024**3):.2f} GB, " # type: ignore[attr-defined] + f"free: {free_memory / (1024**3):.2f} GB, total: {total_memory / (1024**3):.2f} GB" + ) self.hooks_config = self._setup_hooks(hooks_config=hooks_config) From f2ecbc96719acb7d2f0f5814ab179fac83438437 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 29 Jul 2026 08:33:38 +0000 Subject: [PATCH 6/7] [Fix] Stage optimizer states during checkpoint resume A resumed process loads optimizer states before its cold first forward/backward, causing their memory to overlap the compilation peak and starve NCCL allocations. Temporarily offload restored GPU optimizer tensors to CPU, then restore each tensor to its original device at the first optimizer-step boundary. --- tests/engine/test_glm52_moe_train_engine.py | 35 ++++++++++++++++++++- tests/train/test_trainer.py | 6 ++++ xtuner/v1/engine/train_engine.py | 26 +++++++++++++++ xtuner/v1/train/trainer.py | 15 +++++++-- 4 files changed, 79 insertions(+), 3 deletions(-) diff --git a/tests/engine/test_glm52_moe_train_engine.py b/tests/engine/test_glm52_moe_train_engine.py index a10240338..f44c2206d 100644 --- a/tests/engine/test_glm52_moe_train_engine.py +++ b/tests/engine/test_glm52_moe_train_engine.py @@ -9,7 +9,7 @@ test_tilewise_fp8_loss_curve_matches_bf16: tilewise FP8 训练轨迹接近 BF16。 test_tilewise_fp8_ep4_train_step: tilewise FP8 与 EP4 联合训练产生有限 loss。 TestGlm52CheckpointEngine - test_dcp_round_trip_preserves_model_and_optimizer: DCP 往返保留模型与优化器状态。 + test_dcp_round_trip_preserves_model_and_optimizer: DCP 往返保留状态，并支持恢复首步暂存优化器状态。 """ import json @@ -464,6 +464,39 @@ def test_dcp_round_trip_preserves_model_and_optimizer(self): torch.equal(actual, expected), f"optimizer state mismatch: {param_id}.{state_key}", ) + + # A resumed process has a cold first forward/backward. Keep the + # restored optimizer tensors on CPU through that peak, then verify + # the optimizer-step boundary restores every tensor exactly. + optimizer_tensors = [] + for state in restored.optimizer.state.values(): + for state_key, value in state.items(): + if not isinstance(value, torch.Tensor): + continue + local_value = value.to_local() if isinstance(value, DTensor) else value + optimizer_tensors.append( + (state, state_key, local_value.device, local_value.detach().cpu().clone()) + ) + + self.assertTrue(restored.offload_optimizer_until_step()) + for state, state_key, original_device, _ in optimizer_tensors: + value = state[state_key] + local_value = value.to_local() if isinstance(value, DTensor) else value + if original_device.type != "cpu": + self.assertEqual(local_value.device.type, "cpu") + + resumed_seq_ctx = SequenceContext.from_input_ids((input_ids[:, :-1],), device=DEVICE) + resumed_data = {"seq_ctx": resumed_seq_ctx, "shifted_labels": input_ids[:, 1:]} + resumed_loss_ctx = restored.model.build_loss_ctx_batch([resumed_data], sp_mesh=None)[0] + restored.train_step([ModelItem(seq_ctx=resumed_seq_ctx, loss_ctx=resumed_loss_ctx)]) + restored_grad_norm = restored.clip_grad_norm() + restored.step_optimizer(torch.full_like(restored_grad_norm, float("nan"))) + + for state, state_key, original_device, expected in optimizer_tensors: + value = state[state_key] + local_value = value.to_local() if isinstance(value, DTensor) else value + self.assertEqual(local_value.device, original_device) + self.assertTrue(torch.equal(local_value.cpu(), expected)) finally: dist.barrier() if dist.get_rank() == 0: diff --git a/tests/train/test_trainer.py b/tests/train/test_trainer.py index a47ced1a4..5c77dc682 100644 --- a/tests/train/test_trainer.py +++ b/tests/train/test_trainer.py @@ -48,6 +48,7 @@ def __init__(self): self.train_step_calls = 0 self.grad_norm_calls = 0 self.optimizer_step_calls = 0 + self.optimizer_offload_calls = 0 self.model = model = nn.Linear(10, 10) self.optimizer = torch.optim.Adam(model.parameters(), lr=0.001) @@ -89,6 +90,10 @@ def clip_grad_norm(self, do_clip: bool = True, dtype=torch.float32): self.grad_norm_calls += 1 return torch.tensor(1.0) + def offload_optimizer_until_step(self): + self.optimizer_offload_calls += 1 + return True + load_dcp = Mock() def save_dcp(self, weights_dir: Path, save_optimizer: bool = True): @@ -1006,6 +1011,7 @@ class TemporaryState: if DEVICE == "cuda": assert torch.cuda.memory_reserved() < load_artifacts["reserved_memory"] assert "[Checkpoint Resume Memory]" in capfd.readouterr().err + assert trainer._engine.optimizer_offload_calls == 1 # assert trainer._load_checkpoint_cfg.load_dataset is False # assert trainer._load_checkpoint_cfg.load_scheduler is False trainer.fit() diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 61a1f58fe..50e5968ec 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -158,6 +158,7 @@ def __init__( self.has_freeze_params = self.__has_freeze_params() self._async_checkpoint_pg: dist.ProcessGroup | None = None self._async_state_dict_cache: dict[str, Any] | None = None + self._optimizer_state_restore_plan: list[tuple[dict[str, Any], str, torch.device]] = [] def __has_freeze_params(self) -> bool: has_freeze_params = False @@ -276,6 +277,14 @@ def clip_grad_norm(self, do_clip: bool = True, dtype=torch.float32): def step_optimizer(self, grad_norm): """Step the optimizer to update the model parameters.""" + if self._optimizer_state_restore_plan: + # Resume keeps optimizer states on CPU through the cold first + # forward/backward, then restores their original placement here. + for state, key, device in self._optimizer_state_restore_plan: + state[key] = state[key].to(device, non_blocking=True) + DEVICE_MODULE.synchronize() + self._optimizer_state_restore_plan = [] + if torch.isnan(grad_norm) or torch.isinf(grad_norm): log_rank0.warning(f"Gradient norm {grad_norm} is invalid, skipping optimizer step.") self.optimizer.zero_grad() @@ -554,6 +563,23 @@ def put_optimizer_to_device(self, device: torch.device | str) -> bool: DEVICE_MODULE.synchronize() return True + def offload_optimizer_until_step(self) -> bool: + """Keep GPU optimizer states on CPU until the next optimizer step.""" + if getattr(self.optim_cfg, "swap_optimizer", False) or not self.optimizer.state: + return False + + restore_plan = [] + for state in self.optimizer.state.values(): + if isinstance(state, dict): + for key, val in state.items(): + if isinstance(val, torch.Tensor) and val.device.type != "cpu": + restore_plan.append((state, key, val.device)) + + if not restore_plan or not self.put_optimizer_to_device("cpu"): + return False + self._optimizer_state_restore_plan = restore_plan + return True + def _maybe_precompute_float8_dynamic_scale_for_fsdp(self): for model in self.model.modules(): if isinstance(model, BaseModel) and model.float8_handler is not None: diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index f3e2b2f24..98ef5395f 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -744,8 +744,19 @@ def __init__( if self._load_checkpoint_cfg.checkpoint_path is not None: self._load_checkpoint() - # DCP load materializes temporary model and optimizer state. Release - # its cached device memory before the first resumed training step. + # A fresh process has a cold first forward/backward. Keep restored + # optimizer states out of that peak and restore them at optimizer.step. + optimizer_offloaded = False + if self._load_checkpoint_cfg.load_optimizer_states and DEVICE != "cpu": + optimizer_offloaded = self._engine.offload_optimizer_until_step() + if optimizer_offloaded: + self.logger.info( + "[Checkpoint Resume] Optimizer states are temporarily offloaded " + "to CPU until the first optimizer step." + ) + + # Release DCP's temporary state and the allocator cache left by the + # optimizer-state transfer before the first resumed training step. gc.collect() DEVICE_MODULE.empty_cache() if DEVICE != "cpu": From cb9878c5bd710a8f7a2d37584b1e5f0362567a24 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 29 Jul 2026 08:56:36 +0000 Subject: [PATCH 7/7] [Refactor] Configure first-step optimizer offload Expose offload_optimizer_first_step on LoadCheckpointConfig and use the existing put_optimizer_to_device API directly. Trainer moves optimizer state to CPU after checkpoint load and back to the training device after the first gradient reduction, without maintaining a separate engine restore plan. --- tests/engine/test_glm52_moe_train_engine.py | 24 +++++++++++-------- tests/train/test_trainer.py | 10 ++++---- xtuner/v1/engine/train_engine.py | 26 --------------------- xtuner/v1/train/trainer.py | 14 ++++++----- 4 files changed, 28 insertions(+), 46 deletions(-) diff --git a/tests/engine/test_glm52_moe_train_engine.py b/tests/engine/test_glm52_moe_train_engine.py index f44c2206d..88dd9386d 100644 --- a/tests/engine/test_glm52_moe_train_engine.py +++ b/tests/engine/test_glm52_moe_train_engine.py @@ -474,29 +474,33 @@ def test_dcp_round_trip_preserves_model_and_optimizer(self): if not isinstance(value, torch.Tensor): continue local_value = value.to_local() if isinstance(value, DTensor) else value - optimizer_tensors.append( - (state, state_key, local_value.device, local_value.detach().cpu().clone()) - ) + optimizer_tensors.append((state, state_key, local_value.detach().cpu().clone())) - self.assertTrue(restored.offload_optimizer_until_step()) - for state, state_key, original_device, _ in optimizer_tensors: + self.assertTrue(restored.put_optimizer_to_device("cpu")) + for state, state_key, _ in optimizer_tensors: value = state[state_key] local_value = value.to_local() if isinstance(value, DTensor) else value - if original_device.type != "cpu": - self.assertEqual(local_value.device.type, "cpu") + self.assertEqual(local_value.device.type, "cpu") resumed_seq_ctx = SequenceContext.from_input_ids((input_ids[:, :-1],), device=DEVICE) resumed_data = {"seq_ctx": resumed_seq_ctx, "shifted_labels": input_ids[:, 1:]} resumed_loss_ctx = restored.model.build_loss_ctx_batch([resumed_data], sp_mesh=None)[0] restored.train_step([ModelItem(seq_ctx=resumed_seq_ctx, loss_ctx=resumed_loss_ctx)]) restored_grad_norm = restored.clip_grad_norm() - restored.step_optimizer(torch.full_like(restored_grad_norm, float("nan"))) + self.assertTrue(restored.put_optimizer_to_device(DEVICE)) - for state, state_key, original_device, expected in optimizer_tensors: + for state, state_key, expected in optimizer_tensors: value = state[state_key] local_value = value.to_local() if isinstance(value, DTensor) else value - self.assertEqual(local_value.device, original_device) + self.assertEqual(local_value.device.type, str(DEVICE)) self.assertTrue(torch.equal(local_value.cpu(), expected)) + + restored.step_optimizer(restored_grad_norm) + for state in restored.optimizer.state.values(): + for value in state.values(): + if isinstance(value, torch.Tensor): + local_value = value.to_local() if isinstance(value, DTensor) else value + self.assertTrue(torch.isfinite(local_value).all()) finally: dist.barrier() if dist.get_rank() == 0: diff --git a/tests/train/test_trainer.py b/tests/train/test_trainer.py index 5c77dc682..aa2b55980 100644 --- a/tests/train/test_trainer.py +++ b/tests/train/test_trainer.py @@ -48,7 +48,7 @@ def __init__(self): self.train_step_calls = 0 self.grad_norm_calls = 0 self.optimizer_step_calls = 0 - self.optimizer_offload_calls = 0 + self.optimizer_device_calls = [] self.model = model = nn.Linear(10, 10) self.optimizer = torch.optim.Adam(model.parameters(), lr=0.001) @@ -90,8 +90,8 @@ def clip_grad_norm(self, do_clip: bool = True, dtype=torch.float32): self.grad_norm_calls += 1 return torch.tensor(1.0) - def offload_optimizer_until_step(self): - self.optimizer_offload_calls += 1 + def put_optimizer_to_device(self, device): + self.optimizer_device_calls.append(str(device)) return True load_dcp = Mock() @@ -961,6 +961,7 @@ def test_resume_and_load_checkpoint_cfg(tmp_path: Path, capfd): load_optimizer_args=True, load_dataset=False, load_scheduler=False, + offload_optimizer_first_step=True, ) load_artifacts = {} @@ -1011,10 +1012,11 @@ class TemporaryState: if DEVICE == "cuda": assert torch.cuda.memory_reserved() < load_artifacts["reserved_memory"] assert "[Checkpoint Resume Memory]" in capfd.readouterr().err - assert trainer._engine.optimizer_offload_calls == 1 + assert trainer._engine.optimizer_device_calls == ["cpu"] # assert trainer._load_checkpoint_cfg.load_dataset is False # assert trainer._load_checkpoint_cfg.load_scheduler is False trainer.fit() + assert trainer._engine.optimizer_device_calls == ["cpu", str(DEVICE)] # 4. 2nd create: resume train with auto_resume and load_checkpoint_cfg with ( diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 50e5968ec..61a1f58fe 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -158,7 +158,6 @@ def __init__( self.has_freeze_params = self.__has_freeze_params() self._async_checkpoint_pg: dist.ProcessGroup | None = None self._async_state_dict_cache: dict[str, Any] | None = None - self._optimizer_state_restore_plan: list[tuple[dict[str, Any], str, torch.device]] = [] def __has_freeze_params(self) -> bool: has_freeze_params = False @@ -277,14 +276,6 @@ def clip_grad_norm(self, do_clip: bool = True, dtype=torch.float32): def step_optimizer(self, grad_norm): """Step the optimizer to update the model parameters.""" - if self._optimizer_state_restore_plan: - # Resume keeps optimizer states on CPU through the cold first - # forward/backward, then restores their original placement here. - for state, key, device in self._optimizer_state_restore_plan: - state[key] = state[key].to(device, non_blocking=True) - DEVICE_MODULE.synchronize() - self._optimizer_state_restore_plan = [] - if torch.isnan(grad_norm) or torch.isinf(grad_norm): log_rank0.warning(f"Gradient norm {grad_norm} is invalid, skipping optimizer step.") self.optimizer.zero_grad() @@ -563,23 +554,6 @@ def put_optimizer_to_device(self, device: torch.device | str) -> bool: DEVICE_MODULE.synchronize() return True - def offload_optimizer_until_step(self) -> bool: - """Keep GPU optimizer states on CPU until the next optimizer step.""" - if getattr(self.optim_cfg, "swap_optimizer", False) or not self.optimizer.state: - return False - - restore_plan = [] - for state in self.optimizer.state.values(): - if isinstance(state, dict): - for key, val in state.items(): - if isinstance(val, torch.Tensor) and val.device.type != "cpu": - restore_plan.append((state, key, val.device)) - - if not restore_plan or not self.put_optimizer_to_device("cpu"): - return False - self._optimizer_state_restore_plan = restore_plan - return True - def _maybe_precompute_float8_dynamic_scale_for_fsdp(self): for model in self.model.modules(): if isinstance(model, BaseModel) and model.float8_handler is not None: diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index 98ef5395f..bd780a7de 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -387,6 +387,7 @@ class LoadCheckpointConfig(BaseModel): load_optimizer_args: bool = True load_dataset: bool = True load_scheduler: bool = True + offload_optimizer_first_step: bool = False class TrainerConfig(BaseModel): @@ -742,14 +743,12 @@ def __init__( self._checkpoint_interval = None self._snapshot_interval = None + self._offload_optimizer_first_step = False if self._load_checkpoint_cfg.checkpoint_path is not None: self._load_checkpoint() - # A fresh process has a cold first forward/backward. Keep restored - # optimizer states out of that peak and restore them at optimizer.step. - optimizer_offloaded = False - if self._load_checkpoint_cfg.load_optimizer_states and DEVICE != "cpu": - optimizer_offloaded = self._engine.offload_optimizer_until_step() - if optimizer_offloaded: + self._offload_optimizer_first_step = self._load_checkpoint_cfg.offload_optimizer_first_step + if self._offload_optimizer_first_step: + self._engine.put_optimizer_to_device("cpu") self.logger.info( "[Checkpoint Resume] Optimizer states are temporarily offloaded " "to CPU until the first optimizer step." @@ -875,6 +874,9 @@ def fit(self): ) grad_norm = self._engine.clip_grad_norm(do_clip=self._do_clip, dtype=self._grad_norm_dtype) + if self._offload_optimizer_first_step: + self._engine.put_optimizer_to_device(DEVICE) + self._offload_optimizer_first_step = False self._engine.step_optimizer(grad_norm) time_after_train_step = time.time()