Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / base.py: 28%
1122 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""BaseTrainer — composable training skeleton with 13 overridable ``_build_*`` steps.
17Design notes:
18- Composition over inheritance: a trainer holds a ``BaseTrainer`` and calls its
19 13 ``_build_*`` steps in order, overriding or skipping steps as needed.
20- FSDP/AC wrapping iterates ``model.layers`` when the model exposes decoder layers.
21- Parallel composition order is TP → CP → AC → FSDP.
23Subclasses (LLMTrainer, VLMTrainer, ...) follow this pattern: instantiate a
24``BaseTrainer`` and drive its ``_build_*`` methods selectively.
25"""
26import json
27import logging
28import math
29import os
30import random
31from contextlib import nullcontext
32from typing import TYPE_CHECKING, Any, Dict, Optional
34import numpy as np
35import torch
36from torch.utils.data import DistributedSampler
38from hyper_parallel import (
39 get_platform,
40 init_empty_weights,
41 init_process_group,
42 destroy_process_group,
43 hsdp_sync_stream,
44 SkipDTensorDispatch,
45 HSDPModule,
46)
47from hyper_parallel.core.distributed_checkpoint import load as dcp_load
48from hyper_parallel.core.dtensor.dtensor import DTensor
49# ``_resolve_local_tensor`` is the canonical shard resolver used by
50# ``HSDPModule.load_state_dict``; reused (rather than duplicated) to load a
51# checkpoint into a model that holds DTensor params but is not itself an
52# ``HSDPModule`` (pipeline parallelism composed with per-module FSDP).
53from hyper_parallel.core.fully_shard.api import _resolve_local_tensor
54from hyper_parallel.core.fully_shard.hsdp_utils import GroupInfo
55from hyper_parallel.core.utils import clip_grad_norm_
56from hyper_parallel.data import build_dataset
57from hyper_parallel.models.spec.registry import get_spec
58from hyper_parallel.trainer.config import get_vision_parallel_config
59from hyper_parallel.trainer.parallel_dims import ParallelDims
60from hyper_parallel.trainer.utils.loss import count_loss_token, mean_global_loss
61from hyper_parallel.trainer.callbacks.base import (
62 LoggingCallback,
63 CheckpointCallback,
64 SafetensorsExportCallback,
65 EvalCallback,
66 ProfilerCallback,
67 WandbCallback,
68 ProgressCallback,
69 MoEMonitorCallback,
70 TrainingStateMonitorCallback,
71 GradientHealthCallback,
72 GCCallback,
73 TensorBoardCallback,
74 MemoryMonitorCallback,
75)
77if TYPE_CHECKING:
78 # Type-only imports — never executed at runtime, so the platform-agnostic
79 # rule ("no torch/mindspore in trainer code") is preserved. Same pattern
80 # as
81 from torch import nn
82 from torch.optim import Optimizer
83 from torch.optim.lr_scheduler import LRScheduler
84 from torch.utils.data import DataLoader
85 from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
87platform = get_platform()
88logger = logging.getLogger(__name__)
91class TrainerState:
92 """Mutable training state shared across callbacks.
94 Attributes:
95 global_step: Current training step (update count).
96 epoch: Current epoch index.
97 max_steps: Total number of training steps.
98 """
100 def __init__(self, max_steps: int = 0):
101 self.global_step: int = 0
102 self.epoch: int = 0
103 self.max_steps: int = max_steps
104 self.log_history: list = []
105 self.substep_info: Dict[str, Any] = {}
108class BaseTrainer:
109 """Composable training skeleton.
111 Provides 13 ``_build_*`` methods that subclasses can call, override, or skip.
112 The default ``_build_parallelized_model`` applies TP → CP → AC → FSDP by
113 iterating ``model.layers`` — matching hyper's own ``fsdp_demo.py`` style.
115 Args:
116 args: Training configuration (typically parsed from YAML).
117 """
119 # PEP 526 annotations — populated by ``_build_*``; ``None`` until built.
120 model: Optional["nn.Module"] = None
121 optimizer: Optional["Optimizer"] = None
122 lr_scheduler: Optional["LRScheduler"] = None
123 train_dataloader: Optional["DataLoader"] = None
124 mesh: Optional["DeviceMesh"] = None
125 # Pipeline-parallel state — set by ``_build_pipelined_model`` when ``pp>1``.
126 pp_enabled: bool = False
127 pp_schedule: Optional[Any] = None
128 pp_micro_batch_num: int = 1
129 pp_has_first_stage: bool = False
130 pp_has_last_stage: bool = False
131 _pp_tie_embeddings: bool = False
132 _pp_stage_fsdp_sharded: bool = False
134 def __init__(self, args):
135 # Only early-bound fields live here; the rest is built via
136 # ``_build_*`` methods invoked by the subclass.
137 self.args = args
138 self.spec = get_spec(args.model.name)
139 self.state = TrainerState(max_steps=args.train.max_steps)
140 self._pp_stage_modules: list["nn.Module"] = []
141 self._pp_tp_loss_repeats = 1
143 # ------------------------------------------------------------------
144 # 13 overridable _build_* methods
145 # ------------------------------------------------------------------
147 @property
148 def _deterministic(self) -> bool:
149 return bool(self.args.train.debug.deterministic)
151 def _apply_pre_init_deterministic_env(self):
152 """Pin HCCL / PYTHONHASHSEED before ``init_process_group`` boots the backend."""
153 if not self._deterministic:
154 return
155 seed = self.args.train.seed
156 os.environ.setdefault("ASCEND_LAUNCH_BLOCKING", "1")
157 os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
158 os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":16:8")
159 os.environ.setdefault("FLASH_ATTENTION_DETERMINISTIC", "1")
160 os.environ.setdefault("HCCL_DETERMINISTIC", "true")
161 os.environ.setdefault("PYTHONHASHSEED", str(seed))
163 def _parallel_dim_size(self, name: str) -> int:
164 """Return a configured parallel dimension size."""
165 return int(getattr(self.parallel_dims, name, 1) or 1)
167 def _cp_size(self) -> int:
168 """Return configured context-parallel size."""
169 return self._parallel_dim_size("cp")
171 def _share_samples_across_dp(self) -> bool:
172 """Return whether the visual validation path reuses samples across DP."""
173 return get_vision_parallel_config(self.args.model).get(
174 "share_samples_across_dp", False,
175 )
177 def _setup(self):
178 """Step 1: Initialize distributed environment, device mesh, and seed.
180 Calls hyper's own ``init_process_group`` and ``init_device_mesh``.
181 Mesh shape is derived from ``args.parallel`` (dp, tp, cp, pp, ep).
182 """
183 self._apply_pre_init_deterministic_env()
184 backend = self.args.train.comm_backend
185 init_process_group(backend=backend)
187 local_rank = self.args.train.local_rank
188 device_type = platform.device_type() # "npu" or "cuda"
189 # Use platform.device(idx) — backend-agnostic.
190 self.device = platform.device(local_rank)
191 device_handle = platform.get_device_handle(device_type)
192 device_handle.set_device(local_rank)
194 # Build & validate parallel dims in one place (fail-fast).
196 self.parallel_dims = ParallelDims.from_config(
197 self.args.train.accelerator, world_size=platform.get_world_size(),
198 )
199 logger.info_rank0("ParallelDims: %s", self.parallel_dims.summary())
200 # Mixed precision lives in FSDP2's MixedPrecisionPolicy, so a
201 # low-precision run needs a dp_shard axis (size-1 is enough) for the
202 # FSDP wrap to exist — see ``build_mesh``'s force_dp_shard contract.
203 mp_cfg = self.args.train.mixed_precision
204 needs_mp_wrap = bool(
205 mp_cfg.enabled
206 and mp_cfg.param_dtype not in ('float32', 'fp32')
207 )
208 # PP stages carry the dtype policy only through a per-stage FSDP wrap,
209 # which exists only for pure dp_shard sharding (no HSDP, see
210 # ``_resolve_fsdp_mesh``) — reject every PP composition that would
211 # silently run full-precision instead.
212 if (needs_mp_wrap and self.parallel_dims.pp > 1
213 and (self.parallel_dims.dp_shard == 1
214 or self.parallel_dims.dp_replicate > 1)
215 and self._cp_size() == 1):
216 raise ValueError(
217 "mixed_precision with a low-precision param_dtype under PP "
218 "needs an FSDP-wrappable data-parallel axis: the dtype policy "
219 "lives on the per-stage FSDP wrap, which neither pure PP nor "
220 "PP+HSDP provides. Use dp_shard>=2 with dp_replicate=1, or "
221 "set param_dtype=float32."
222 )
223 self.mesh = self.parallel_dims.build_mesh(
224 platform.device_type(), force_dp_shard=needs_mp_wrap,
225 )
227 # Build DP group_info for trainer-level all_reduce (loss/token sync).
228 # Uses hyper's GroupInfo + mesh.get_group (platform-agnostic).
230 dp_group = self._get_combined_dp_group()
231 dp_size = self.parallel_dims.dp_size
232 self._dp_group_info = GroupInfo(
233 group_name="trainer_dp", group=dp_group, rank_size=dp_size,
234 )
236 seed = self.args.train.seed
237 platform.manual_seed(seed)
238 random.seed(seed)
239 np.random.seed(seed)
240 # ``platform.manual_seed`` only covers CPU; seed the device RNG too.
241 try:
242 handle = platform.get_device_handle(device_type)
243 if hasattr(handle, "manual_seed_all"):
244 handle.manual_seed_all(seed)
245 elif hasattr(handle, "manual_seed"):
246 handle.manual_seed(seed)
247 except Exception as exc: # pylint: disable=W0718
248 logger.warning("Device-side seed init skipped: %s", exc)
250 if self._deterministic:
251 warn_only = self.args.train.debug.deterministic_warn_only
252 torch.use_deterministic_algorithms(True, warn_only=warn_only)
253 torch.backends.cudnn.deterministic = True
254 torch.backends.cudnn.benchmark = False
255 # TF32 affects CUDA only; the attribute may be missing on older torch.
256 try:
257 torch.backends.cuda.matmul.allow_tf32 = False
258 torch.backends.cudnn.allow_tf32 = False
259 except AttributeError:
260 pass
261 logger.info_rank0("Deterministic algorithms enabled (warn_only=%s)", warn_only)
263 logger.info_rank0(
264 "Setup complete: rank=%d, world_size=%d, mesh=%s",
265 platform.get_rank(), platform.get_world_size(),
266 self.mesh.mesh_dim_names,
267 )
268 logger.info_rank0(
269 "Config: data.type=%s, model.name=%s, model.num_hidden_layers=%s, "
270 "init_device=%s, max_steps=%d, global_bs=%d",
271 self.args.data.type,
272 self.args.model.name,
273 self.args.model.num_hidden_layers,
274 self.args.train.init_device,
275 self.state.max_steps,
276 self.args.train.global_batch_size,
277 )
279 def _build_model(self):
280 """Step 2: Construct model via ``spec.build_model_fn``.
282 The model is a plain ``nn.Module`` at this point — not yet parallelized.
283 When ``args.runtime.init_device == "meta"``, the model is constructed on
284 the meta device (no memory allocated) and real weights are loaded after
285 FSDP sharding via ``_load_weights_after_parallel``.
286 """
287 init_device = self.args.train.init_device
288 # Meta-device init: each rank materialises only its own shard
289 # post-FSDP — pre-trained weights via DCP, otherwise random init.
290 if init_device == "meta":
292 with init_empty_weights():
293 self.model = self.spec.build_model_fn(self.args)
294 logger.info_rank0(
295 "Model built on meta device (no memory allocated): %s",
296 type(self.model).__name__,
297 )
298 else:
299 self.model = self.spec.build_model_fn(self.args)
300 logger.info_rank0("Model built on %s: %s", init_device, type(self.model).__name__)
302 # Cross-check parallel degrees against the actual model hyperparams
303 # (heads%tp, kv_heads%tp, num_experts%ep, seq_len%(cp*tp)).
304 # Fails fast here instead of crashing inside parallelize_module.
305 seq_len = self.args.data.max_seq_len
306 self.parallel_dims.validate_against_model(self.model, seq_len=seq_len)
308 def _freeze_model(self):
309 """Step 3: Freeze specified modules (optional)."""
310 freeze_modules = self.args.model.freeze_modules
311 if not freeze_modules:
312 return
313 for name, param in self.model.named_parameters():
314 if any(pattern in name for pattern in freeze_modules):
315 param.requires_grad_(False)
317 def _build_model_assets(self):
318 """Step 4: Build tokenizer, processor, chat_template.
320 Default: no-op. LLMTrainer overrides to build tokenizer + chat_template.
321 VLMTrainer overrides to build processor.
322 """
323 self.tokenizer = None
324 self.processor = None
326 def _build_data_transform(self):
327 """Step 5: Build data preprocessing transform.
329 Default: identity transform. LLMTrainer overrides for tokenization.
330 """
331 self.data_transform = None
333 def _build_dataset(self):
334 """Step 6: Build training dataset via the data-type registry.
336 Dispatches on ``args.data.type`` against
337 :data:`hyper_parallel.data.DATASET_REGISTRY`. Built-in formats:
338 ``dummy``, ``hf_datasets``, ``json_file``, ``preset_pt``,
339 ``vl_dummy``, ``megatron``. Plug in a custom format by importing
340 a module that calls ``@DATASET_REGISTRY.register(...)``.
342 Subclasses can override to populate ``self.train_dataset``
343 differently before this method runs (or skip it entirely).
344 """
345 if getattr(self, "train_dataset", None) is not None:
346 return
347 if self.args.data.streaming:
348 # ``DistributedSampler`` requires ``__len__``; an iterable path
349 # would need a sampler-less dataloader. Reject loudly until that
350 # path is wired so users see a clear error instead of a
351 # ``TypeError: object of type ... has no len()``.
352 raise NotImplementedError(
353 "data.streaming=True is not yet wired. The default "
354 "_build_dataloader uses DistributedSampler which requires "
355 "len(dataset); subclass _build_dataset + _build_dataloader "
356 "to emit an IterableDataset that self-shards via dp_rank/dp_size."
357 )
358 data_type = self.args.data.type
359 self.train_dataset = build_dataset(
360 data_type,
361 base=self,
362 args=self.args,
363 tokenizer=getattr(self, "tokenizer", None),
364 data_transform=getattr(self, "data_transform", None),
365 )
367 def _build_collate_fn(self):
368 """Step 7: Build data collator.
370 Default: pads input_ids and labels to max length in the batch.
371 SequenceParallel TP and context parallel both slice the sequence
372 dim, so variable-length batches additionally pad up to a multiple
373 of ``cp * tp`` — the trailing pad carries label ``-100``, which the
374 CE masks out, so the padding is mathematically inert.
375 """
376 seq_divisor = self.parallel_dims.seq_divisor
378 def _default_collate(batch):
379 """Simple padding collator."""
380 max_len = max(item["input_ids"].size(0) for item in batch)
381 if seq_divisor > 1 and max_len % seq_divisor:
382 max_len += seq_divisor - max_len % seq_divisor
383 input_ids_list = []
384 labels_list = []
385 for item in batch:
386 pad_len = max_len - item["input_ids"].size(0)
387 input_ids_list.append(
388 torch.nn.functional.pad(item["input_ids"], (0, pad_len), value=0)
389 )
390 labels_list.append(
391 torch.nn.functional.pad(item["labels"], (0, pad_len), value=-100)
392 )
393 out = {
394 "input_ids": torch.stack(input_ids_list),
395 "labels": torch.stack(labels_list),
396 }
397 if "num_items_in_batch" in batch[0]:
398 out["num_items_in_batch"] = sum(
399 int(item["num_items_in_batch"]) for item in batch
400 )
401 if "attention_mask" in batch[0]:
402 masks = []
403 for item in batch:
404 pad_len = max_len - item["attention_mask"].size(0)
405 masks.append(torch.nn.functional.pad(item["attention_mask"], (0, pad_len), value=0))
406 out["attention_mask"] = torch.stack(masks)
407 if "position_ids" in batch[0]:
408 positions = []
409 for item in batch:
410 pos = item["position_ids"]
411 pad_len = max_len - pos.shape[-1]
412 positions.append(torch.nn.functional.pad(pos, (0, pad_len), value=0))
413 if positions[0].dim() == 1:
414 out["position_ids"] = torch.stack(positions)
415 else:
416 out["position_ids"] = torch.stack(positions).transpose(0, 1).contiguous()
417 return out
419 self.collate_fn = _default_collate
421 def _build_dataloader(self):
422 """Step 8: Build distributed stateful dataloader.
424 Uses ``torchdata.stateful_dataloader.StatefulDataLoader`` so that
425 iterator position is checkpointable — enabling exact resume after
426 restart (matching ).
428 Each ``next()`` call yields a list of micro-batches (for gradient
429 accumulation).
430 """
431 from torchdata.stateful_dataloader import StatefulDataLoader # pylint: disable=C0415 # optional dep
433 micro_bs = self.args.train.micro_batch_size
435 # Sampler uses DP rank/size — TP/CP/PP/EP peers share data.
436 dp_size = self.parallel_dims.dp_size
437 non_dp = self.parallel_dims.non_dp_size
438 global_rank = platform.get_rank()
439 try:
440 dp_rank = self.mesh["dp"].get_local_rank()
441 except (KeyError, ValueError, RuntimeError):
442 dp_rank = global_rank // non_dp if non_dp > 1 else global_rank
444 shuffle = self.args.data.shuffle
445 sampler_seed = self.args.train.seed
446 self.sampler = DistributedSampler(
447 self.train_dataset,
448 num_replicas=1 if self._share_samples_across_dp() else dp_size,
449 rank=0 if self._share_samples_across_dp() else dp_rank,
450 shuffle=shuffle,
451 seed=sampler_seed,
452 drop_last=True,
453 )
455 # StatefulDataLoader supports state_dict() / load_state_dict()
456 # for checkpoint resume (torchdata API, used by + ).
457 num_workers = self.args.data.num_workers
458 prefetch_factor = self.args.data.prefetch_factor
459 pin_memory = self.args.data.pin_memory
461 # Spawned-worker RNG is not bit-stable across 1c↔Nc; force num_workers=0
462 # in deterministic mode.
463 if self._deterministic and num_workers > 0:
464 logger.warning(
465 "debug.deterministic=True forces data.num_workers from %d → 0",
466 num_workers,
467 )
468 num_workers = 0
470 loader_kwargs = {
471 "batch_size": micro_bs,
472 "sampler": self.sampler,
473 "collate_fn": self.collate_fn,
474 "num_workers": num_workers,
475 "pin_memory": pin_memory,
476 "drop_last": True,
477 }
478 # prefetch_factor is only accepted when num_workers > 0
479 if num_workers > 0 and prefetch_factor is not None:
480 loader_kwargs["prefetch_factor"] = prefetch_factor
481 if self._deterministic:
482 # Pin loader RNG to the trainer seed so shuffle order is stable.
483 gen = torch.Generator()
484 gen.manual_seed(int(self.args.train.seed))
485 loader_kwargs["generator"] = gen
486 self.train_dataloader = StatefulDataLoader(
487 self.train_dataset, **loader_kwargs,
488 )
490 # Use dp_size (not world_size) — TP/CP/PP ranks share data, not split it.
491 self._grad_accum = max(
492 self.args.train.global_batch_size // (
493 micro_bs * (1 if self._share_samples_across_dp() else dp_size)
494 ),
495 1,
496 )
498 if self._share_samples_across_dp():
499 logger.warning_rank0(
500 "vision_parallel.share_samples_across_dp=true. Use this only for "
501 "validation/self-consistency checks; normal training should keep "
502 "distinct samples across DP ranks."
503 )
504 logger.info_rank0(
505 "Dataloader built: micro_bs=%d, grad_accum=%d, dataset_size=%d, "
506 "share_samples_across_dp=%s",
507 micro_bs, self._grad_accum, len(self.train_dataset),
508 str(self._share_samples_across_dp()),
509 )
511 def _build_parallelized_model(self):
512 """Step 9: Apply parallel strategies to the model.
514 Each model owns its full parallelize pipeline in
515 ``models/<name>/parallelize.py`` (convention) and
516 registers it via ``ModelSpec.parallelize_fn``. There is no shared
517 "default" template — model-specific TP/EP/CP/AC/FSDP/Prefetch
518 composition lives next to the model that needs it.
519 """
520 if self.parallel_dims.pp_enabled:
521 self._build_pipelined_model()
522 return
523 if self.spec.parallelize_fn is None:
524 raise ValueError(
525 f"Model '{self.spec.name}' has no ``parallelize_fn`` registered "
526 f"on its ModelSpec. Each model must own its parallelize "
527 f"pipeline in models/<name>/parallelize.py."
528 )
529 self.model = self.spec.parallelize_fn(self.model, self.mesh, self.args)
530 self._post_parallelize()
532 def _validate_pp_model_parallel_grad_clipping(self, dims) -> None:
533 """Reject PP model-parallel clipping until DTensor norms are placement-aware."""
534 max_grad_norm = float(self.args.train.optimizer.max_grad_norm)
535 if max_grad_norm > 0 and (dims.tp > 1 or dims.ep > 1):
536 raise NotImplementedError(
537 "Trainer PP with TP or EP requires max_grad_norm=0: the current "
538 "pipeline gradient norm does not yet deduplicate replicated "
539 "DTensor placements while reducing TP/EP shards."
540 )
542 def _set_pp_stage_modules(self, stages: list[Any]) -> None:
543 """Expose local stage modules and validate their data-parallel representation."""
544 if len(stages) == 1:
545 self.model = stages[0].submodule
546 else:
547 self.model = torch.nn.ModuleList([stage.submodule for stage in stages])
548 self._pp_stage_fsdp_sharded = any(
549 isinstance(module, HSDPModule)
550 for module in self.model.modules()
551 )
552 has_plain_dtensor = any(isinstance(param, DTensor) for param in self.model.parameters())
553 if self._pp_fsdp_composed and not self._pp_stage_fsdp_sharded and has_plain_dtensor:
554 raise NotImplementedError(
555 "Trainer PP data-parallel fallback cannot synchronize DTensor "
556 "stage parameters across the combined DP group. Use dp_shard "
557 "with dp_replicate=1 for PP+TP/EP, or disable TP/EP when using "
558 "PP with dp_replicate>1."
559 )
561 def _validate_pp_runtime_options(self, dims) -> int:
562 """Validate PP loss, batch, checkpointing, and export options."""
563 # The PP loss/grad is normalized to the global token mean. This is
564 # equivalent to ``rank_average`` when every row has the same number of
565 # valid labels; the runtime validates that case before scheduling.
566 agg = self.args.train.optimizer.loss_aggregation
567 if agg not in ('token_weighted', 'rank_average'):
568 raise NotImplementedError(
569 f"Trainer PP supports loss_aggregation='token_weighted' or "
570 f"'rank_average' with uniform valid-token rows only (got {agg!r})."
571 )
573 # The schedule sees the effective batch after the dataloader floors the
574 # configured global batch, so validate the effective size here.
575 micro_num = int(self.args.train.accelerator.pp_micro_batch_num)
576 if micro_num < 1:
577 raise ValueError(f"pp_micro_batch_num ({micro_num}) must be >= 1.")
578 global_bs = self.args.train.global_batch_size
579 micro_bs = int(self.args.train.micro_batch_size)
580 grad_accum = max(int(global_bs) // (micro_bs * dims.dp_size), 1)
581 effective_bs = grad_accum * micro_bs
582 if effective_bs % micro_num != 0:
583 raise ValueError(
584 f"effective PP batch ({effective_bs} = grad_accum*"
585 f"micro_batch_size, floored from global_batch_size={global_bs}) "
586 f"must be divisible by pp_micro_batch_num ({micro_num}); "
587 f"adjust global_batch_size / micro_batch_size / pp_micro_batch_num."
588 )
590 # The PP path bypasses ``parallelize_fn`` and replaces the full model
591 # with a stage fragment, so AC and HF-weight export are not yet wired.
592 ac_mode = self.args.train.gradient_checkpointing.activation_checkpoint
593 if ac_mode not in ("off", "none", None, False, ""):
594 raise NotImplementedError(
595 f"activation_checkpoint={ac_mode!r} is not yet wired for the "
596 f"trainer PP path; set gradient_checkpointing.activation_checkpoint "
597 f"to 'none' for pp>1."
598 )
599 if self.args.train.checkpoint.save_hf_weights:
600 raise NotImplementedError(
601 "checkpoint.save_hf_weights is not yet supported under the "
602 "trainer PP path (each rank holds only a stage fragment); set "
603 "save_hf_weights=false for pp>1."
604 )
605 return micro_num
607 def _build_pipelined_model(self) -> None:
608 """Pipeline-parallel build path (``pp > 1``).
610 Unlike the ``parallelize_fn`` path, the model is **first** materialized
611 and weight-loaded as the *full* network (``_post_parallelize`` is FSDP-
612 agnostic — ``to_empty`` + ``load_state_dict(strict=False)`` work on an
613 unwrapped module), then handed to ``spec.pipelining_fn`` which slices it
614 into this rank's :class:`Qwen3_5StageModule` and returns the
615 ``ScheduleGPipe`` + stages. ``self.model`` is then re-pointed at the
616 stage module so the optimizer / grad-clip built next see only this
617 rank's stage parameters.
619 The trainer supports PP alone and the model-provided FSDP/TP/EP
620 compositions validated below. Unsupported domains such as PP+CP,
621 model-parallel clipping without a placement-aware norm, and plain-DP
622 fallback over DTensor stage parameters fail before training starts.
623 """
624 if self.spec.pipelining_fn is None:
625 raise ValueError(
626 f"Model '{self.spec.name}' has parallel.pp>1 but no "
627 f"``pipelining_fn`` registered on its ModelSpec. Register the "
628 f"model's pipeline splitter (e.g. ``pipeline_<name>_for_trainer``)."
629 )
630 dims = self.parallel_dims
631 # PP composed with FSDP (dp_shard / dp_replicate): each stage's children
632 # are wrapped as FSDP units (load-before-shard) and the 1F1B schedule
633 # defers grad reduction to the final micro-batch backward — every micro
634 # accumulates the unsharded grad locally, then the explicit
635 # FSDP_REDUCE_GRAD step reduces once (see the torch pipeline stage's
636 # per-micro grad-sync defer + ``PipelineStage.execute_reduce_grad``).
637 # EP shards experts within each layer (intra-stage). TP / CP shard the
638 # token sequence; the pipeline carries the sequence-sharded hidden states
639 # across stages (lm_head re-gathers for a full-sequence loss).
640 if dims.cp > 1:
641 raise NotImplementedError(
642 "Trainer pipeline parallelism supports PP alone, PP+FSDP, "
643 f"PP+EP+FSDP, or PP+TP+FSDP (got cp={dims.cp}). Composing PP with "
644 "CP is not yet wired."
645 )
646 self._validate_pp_model_parallel_grad_clipping(dims)
647 self._pp_fsdp_composed = dims.dp_shard > 1 or dims.dp_replicate > 1
648 micro_num = self._validate_pp_runtime_options(dims)
649 # Capture the tie flag while ``self.model`` is still the full model — the
650 # PP grad-clip dedups the tied embed / lm_head, which otherwise lives on
651 # two stages (stage 0's ``embed_tokens`` + the last stage's ``lm_head``).
652 self._pp_tie_embeddings = bool(
653 getattr(self.model.config, "tie_word_embeddings", False)
654 )
655 init_device = self.args.train.init_device
656 if self._pp_fsdp_composed:
657 if init_device != "meta":
658 raise NotImplementedError(
659 "Trainer PP+FSDP currently requires init_device='meta' "
660 f"(got {init_device!r}): each stage's FSDP units are sharded "
661 "on the meta device, then materialized + weight-loaded as "
662 "shards — the same meta path as non-PP FSDP."
663 )
664 # Wrap-on-meta then materialize: ``pipelining_fn`` splits the meta
665 # model and ``fully_shard``-wraps the stage's children, producing
666 # correctly-sized meta shards. ``_post_parallelize`` then runs while
667 # ``self.model`` is still the full model, so ``_load_weights`` maps
668 # the checkpoint by the full-model parameter names (the stage shares
669 # those exact param objects, so its shards receive the weights too).
670 # Doing it the other way round (materialize full → ``fully_shard`` a
671 # real param) leaves the loaded full tensor in place and trips FSDP's
672 # sharded-size check at the first forward.
673 self.pp_schedule, stages = self.spec.pipelining_fn(
674 self.model, self.mesh, self.args,
675 )
676 self._pp_stage_modules = [stage.submodule for stage in stages]
677 self._post_parallelize()
678 # The stage was built while the model was still on meta (so
679 # ``fully_shard`` could create meta shards), which left
680 # ``stage.device`` on meta. ``_post_parallelize`` materialized the
681 # params to the real device; point the stage there too so its P2P
682 # activation buffers — allocated lazily on ``stage.device`` — land
683 # on the compute device instead of meta.
684 for stage in stages:
685 stage.device = self.device
686 # The stage's init-time shared-parameter broadcast was skipped on
687 # meta; now that the shards are materialized + weight-loaded, sync
688 # the tied embed / lm_head ends so both stages start identical.
689 stage._sync_shared_parameters() # pylint: disable=protected-access
690 else:
691 # PP alone: materialize + load the full model, then split (no FSDP
692 # wrap). The full model must be on the trainer device before the
693 # split so a CPU ``init_device`` doesn't leave stages on CPU while
694 # ``_pp_train_step`` moves batches to ``self.device``.
695 self._post_parallelize()
696 self.model = self.model.to(self.device)
697 self.pp_schedule, stages = self.spec.pipelining_fn(
698 self.model, self.mesh, self.args,
699 )
700 self._pp_stage_modules = [stage.submodule for stage in stages]
701 self._pp_tp_loss_repeats = max(int(getattr(self.model, "hp_loss_tp_scale_size", 1)), 1)
702 pp_mesh = self.mesh["pp"]
703 pp_rank = pp_mesh.get_local_rank()
704 self.pp_enabled = True
705 self.pp_micro_batch_num = micro_num
706 self.pp_has_first_stage = pp_rank == 0
707 self.pp_has_last_stage = pp_rank == pp_mesh.size() - 1
708 # Pipeline group for broadcasting the last stage's loss to every rank.
709 self._pp_group_info = GroupInfo(
710 group_name="trainer_pp", group=pp_mesh.get_group(),
711 rank_size=pp_mesh.size(),
712 )
713 # First stage's global rank — the broadcast source for single-reader
714 # data loading in ``_pp_train_step`` (constant, so resolve it once).
715 self._pp_src_rank = platform.get_global_rank(pp_mesh.get_group(), 0)
716 # Re-point ``self.model`` at this rank's stage(s) so the optimizer and
717 # gradient clipping operate on the stage parameters only. Under VPP a
718 # rank owns several non-contiguous chunks; expose all their submodules
719 # (a ModuleList) so every chunk's params are optimized / clipped.
720 self._set_pp_stage_modules(stages)
721 logger.info_rank0(
722 "Pipeline build: pp_size=%d, this rank is stage %d (first=%s, last=%s)",
723 pp_mesh.size(), pp_rank, self.pp_has_first_stage, self.pp_has_last_stage,
724 )
726 def _post_parallelize(self):
727 """Common steps after parallelization (materialize weights + train mode).
729 Order when ``init_device == "meta"`` and ``weights_path`` is set:
731 1. Run ``_materialize_and_init_shards`` first — this calls
732 ``model.to_empty(device=...)`` + kaiming / zero init for every
733 parameter. That is the **baseline** state so no param stays on
734 meta (which would trip ``HSDPState._validate_no_meta_params``).
735 2. Then ``_load_weights`` copies the upstream checkpoint on top.
736 Every key that matches overwrites the random init; anything
737 missing in the checkpoint stays with its kaiming / zero init.
739 This pattern handles partial checkpoints cleanly: any parameter the
740 checkpoint does not supply (e.g. a reduced-layer run where the loader
741 filters out higher layers' keys) keeps its kaiming / zero init, while
742 every key the checkpoint does provide overwrites it. The full Qwen3-VL-
743 MoE checkpoint supplies every module the model defines — ``q_norm`` /
744 ``k_norm`` (per text layer), the vision ``pos_embed`` and
745 ``deepstack_merger_list`` included — so a complete load leaves nothing
746 random.
747 """
748 init_device = self.args.train.init_device
749 weights_path = self.args.model.weights_path
750 if init_device == "meta":
751 # Always materialize first (random init baseline) so no param
752 # stays on meta — then overlay the checkpoint.
753 self._materialize_and_init_shards()
754 if weights_path:
755 self._load_weights(weights_path)
756 elif weights_path:
757 self._load_weights(weights_path)
758 # Mixed-precision storage policy: respect the configured param_dtype
759 # for both trainable and frozen params so optimizer state follows the
760 # same precision contract the forward advertises.
761 self._maybe_downcast_frozen_params()
762 self._maybe_cast_trainable_params()
763 self.model.train()
765 def _maybe_downcast_frozen_params(self) -> None:
766 """Maybe downcast frozen params (internal)."""
767 freeze_modules = self.args.model.freeze_modules
768 if not freeze_modules:
769 return
770 mp_cfg = self.args.train.mixed_precision
771 if not mp_cfg.enabled:
772 return
774 target_dtype = {
775 'bfloat16': torch.bfloat16,
776 'bf16': torch.bfloat16,
777 'float16': torch.float16,
778 'fp16': torch.float16,
779 }.get(mp_cfg.param_dtype)
780 if target_dtype is None:
781 return
782 n_cast = 0
783 for name, param in self.model.named_parameters():
784 if not any(pat in name for pat in freeze_modules):
785 continue
786 if param.requires_grad:
787 continue
788 local = param.data
789 if hasattr(local, 'to_local'):
790 local = local.to_local()
791 if local.dtype == target_dtype:
792 continue
793 new_local = local.to(target_dtype)
794 # DTensor: rebuild the global view via from_local with same placements.
795 if hasattr(param.data, 'to_local'):
796 if isinstance(param.data, DTensor):
797 param.data = DTensor.from_local(
798 new_local,
799 device_mesh=param.data.device_mesh,
800 placements=param.data.placements,
801 )
802 else:
803 param.data = new_local
804 else:
805 param.data = new_local
806 n_cast += 1
807 logger.info_rank0(
808 "Post-load: cast %d frozen params to %s",
809 n_cast, target_dtype,
810 )
812 def _maybe_cast_trainable_params(self) -> None:
813 """Cast trainable params to the configured mixed-precision storage dtype."""
814 mp_cfg = self.args.train.mixed_precision
815 if not mp_cfg.enabled:
816 return
818 dtype_map = {
819 'bfloat16': torch.bfloat16,
820 'bf16': torch.bfloat16,
821 'float16': torch.float16,
822 'fp16': torch.float16,
823 'float32': torch.float32,
824 'fp32': torch.float32,
825 }
826 target_dtype = dtype_map.get(mp_cfg.param_dtype)
827 if target_dtype is None:
828 return
829 target_reduce_dtype = dtype_map.get(mp_cfg.reduce_dtype)
831 def _get_param_local_tensor(param: platform.Parameter) -> platform.Tensor:
832 data = param.data
833 if isinstance(data, DTensor):
834 return data.to_local()
835 return data
837 def _set_param_local_tensor(param: platform.Parameter, local: platform.Tensor) -> None:
838 data = param.data
839 if isinstance(data, DTensor):
840 param.data = DTensor.from_local(
841 local,
842 device_mesh=data.device_mesh,
843 placements=data.placements,
844 shape=tuple(data.shape),
845 )
846 else:
847 param.data = local
849 def _cast_param_data(param: platform.Parameter) -> bool:
850 if not param.requires_grad:
851 return False
852 local = _get_param_local_tensor(param)
853 if local.dtype == target_dtype:
854 return False
855 new_local = local.to(target_dtype)
856 _set_param_local_tensor(param, new_local)
857 return True
859 n_cast = 0
860 seen_param_ids = set()
861 for _, param in self.model.named_parameters():
862 seen_param_ids.add(id(param))
863 if _cast_param_data(param):
864 n_cast += 1
865 def _refresh_hsdp_dtype(hsdp_param) -> None:
866 hsdp_param.orig_dtype = target_dtype
867 hsdp_param.param_dtype = None
868 hsdp_param.reduce_dtype = (
869 None if target_reduce_dtype == target_dtype else target_reduce_dtype
870 )
871 hsdp_param.unsharded_param_buffers = []
872 hsdp_param.reset_sharded_param()
873 if hasattr(hsdp_param, "_unsharded_param"):
874 delattr(hsdp_param, "_unsharded_param")
876 def _refresh_hsdp_state_dtype(state) -> None:
877 if state.param_group is None:
878 return
879 state.param_group.reset_iter_state()
880 state.param_group.all_gather_buckets = []
882 for state in self._iter_hsdp_states():
883 buckets = (
884 getattr(state, 'replicate_params', []) or [],
885 getattr(state, 'hsdp_params', []) or [],
886 )
887 for bucket in buckets:
888 for hsdp_param in bucket:
889 param = getattr(hsdp_param, 'sharded_param', None)
890 if param is None:
891 continue
892 if id(param) not in seen_param_ids and _cast_param_data(param):
893 n_cast += 1
894 seen_param_ids.add(id(param))
895 _refresh_hsdp_dtype(hsdp_param)
896 _refresh_hsdp_state_dtype(state)
897 logger.info_rank0(
898 "Post-load: cast %d trainable params to %s", n_cast, target_dtype,
899 )
901 def _build_optimizer(self):
902 """Step 10: Build optimizer. Must be called AFTER ``_build_parallelized_model``.
904 After FSDP, parameters are DTensor shards — optimizer operates on local shards.
905 Optimizer must be created after ``fully_shard``.
906 """
907 lr = self.args.train.optimizer.lr
908 weight_decay = self.args.train.optimizer.weight_decay
910 # bias / LayerNorm / RMSNorm go to no-decay; grouping matters even
911 # at wd=0 — foreach Adam reduction order differs per group on NPU.
912 decay_keywords = ("bias", "layernorm", "norm", "rmsnorm")
914 def _is_no_decay(name: str) -> bool:
915 lname = name.lower()
916 return any(kw in lname for kw in decay_keywords)
918 decay_params = []
919 no_decay_params = []
920 seen_ids = set()
921 for n, p in self.model.named_parameters():
922 if not p.requires_grad:
923 continue
924 # Dedup tied params (same nn.Parameter shared across modules).
925 if id(p) in seen_ids:
926 continue
927 seen_ids.add(id(p))
928 if _is_no_decay(n):
929 no_decay_params.append(p)
930 else:
931 decay_params.append(p)
933 param_groups = [
934 {"params": decay_params, "weight_decay": weight_decay},
935 {"params": no_decay_params, "weight_decay": 0.0},
936 ]
937 adam_eps = self.args.train.optimizer.eps
938 adam_betas = self.args.train.optimizer.betas
939 adam_foreach = self.args.train.optimizer.foreach
940 # ``None`` intentionally follows PyTorch/HF ``adamw_torch`` defaults.
941 # Deterministic mode controls algorithm selection globally; it should not
942 # silently change the optimizer kernel unless the YAML asks for it.
943 self.optimizer = torch.optim.AdamW(
944 param_groups,
945 lr=lr,
946 betas=adam_betas,
947 eps=adam_eps,
948 foreach=adam_foreach,
949 )
950 logger.info_rank0(
951 "Optimizer: AdamW lr=%.2e wd=%.3g decay_params=%d no_decay_params=%d",
952 lr, weight_decay, len(decay_params), len(no_decay_params),
953 )
955 def _build_lr_scheduler(self):
956 """Step 11: Build learning rate scheduler.
958 Supports cosine decay with warmup. Falls back to constant LR if
959 warmup_ratio is 0 and decay_style is 'constant'.
960 """
962 total_steps = self.state.max_steps
963 warmup_ratio = self.args.train.optimizer.lr_warmup_ratio
964 # ``ceil`` matches the standard warmup convention so a fractional
965 # ``warmup_ratio * max_steps`` rounds up to the next full step.
966 warmup_steps = math.ceil(total_steps * warmup_ratio)
967 decay_style = self.args.train.optimizer.lr_decay_style
968 lr_min = self.args.train.optimizer.lr_min
969 lr_max = self.args.train.optimizer.lr
971 def _lr_lambda(current_step):
972 if current_step < warmup_steps:
973 return float(current_step) / float(max(1, warmup_steps))
974 if decay_style == 'constant':
975 return 1.0
976 # Cosine decay
977 progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
978 cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
979 min_ratio = lr_min / lr_max if lr_max > 0 else 0.0
980 return min_ratio + (1.0 - min_ratio) * cosine_decay
982 self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, _lr_lambda)
983 logger.info_rank0(
984 "LR scheduler: %s, warmup_steps=%d/%d, lr=%.2e→%.2e",
985 decay_style, warmup_steps, total_steps, lr_max, lr_min,
986 )
988 def _build_training_context(self):
989 """Step 12: Build forward/backward context managers.
991 Mixed precision is realised entirely through FSDP2
992 ``MixedPrecisionPolicy`` (param_dtype / reduce_dtype / output_dtype).
993 No autocast context is entered — the model's own ``.float()`` /
994 ``.to(weight.dtype)`` cast points handle the fp32 residual stream.
995 """
996 mp_cfg = self.args.train.mixed_precision
997 self.model_fwd_context = nullcontext()
998 self.model_bwd_context = nullcontext()
999 self.grad_scaler = None
1000 if mp_cfg.enabled:
1001 logger.info_rank0(
1002 "Mixed precision via FSDP2 mp_policy: param=%s reduce=%s on %s",
1003 mp_cfg.param_dtype,
1004 mp_cfg.reduce_dtype,
1005 platform.device_type(),
1006 )
1008 def _init_callbacks(self):
1009 """Step 13: Initialize callbacks (explicit mode).
1011 Each callback is a named field — engineer sees all callbacks and their
1012 order in ``on_step_end`` at a glance. Add/remove/reorder = change one line.
1013 """
1014 self.logging_callback = LoggingCallback(self)
1015 self.checkpoint_callback = CheckpointCallback(self)
1016 self.hf_export_callback = SafetensorsExportCallback(self)
1017 self.eval_callback = EvalCallback(self)
1018 self.profiler_callback = ProfilerCallback(self)
1019 self.wandb_callback = WandbCallback(self)
1020 self.tensorboard_callback = TensorBoardCallback(self)
1021 self.progress_callback = ProgressCallback(self)
1022 self.moe_monitor_callback = MoEMonitorCallback(self)
1023 # Health + operability (no-ops unless enabled in cfg.train.debug / .memory_monitor).
1024 self.training_state_monitor_callback = TrainingStateMonitorCallback(self)
1025 self.gradient_health_callback = GradientHealthCallback(self)
1026 self.memory_monitor_callback = MemoryMonitorCallback(self)
1027 self.gc_callback = GCCallback(self)
1028 # ``user_callbacks`` lets external code append extra Callback instances
1029 # (e.g. domain-specific monitors) without editing this method. They get
1030 # the same lifecycle dispatch as built-ins.
1031 self.user_callbacks: list = []
1032 logger.info_rank0(
1033 "Callbacks initialized: logging, checkpoint, hf_export, eval, "
1034 "profiler, wandb, tensorboard, progress, moe_monitor, "
1035 "training_state_monitor, "
1036 "gradient_health, memory_monitor, gc"
1037 )
1039 # ------------------------------------------------------------------
1040 # Public API: external callback registration
1041 # ------------------------------------------------------------------
1043 def add_callback(self, callback) -> None:
1044 """Register an extra ``Callback`` to receive every lifecycle event.
1046 Use this to plug domain-specific monitors (custom metric sinks,
1047 in-house experiment trackers, RL reward loggers) without editing
1048 the trainer. Built-in callbacks always run first; user callbacks
1049 run in registration order so a later user callback can read state
1050 the earlier ones updated.
1051 """
1052 self.user_callbacks.append(callback)
1053 logger.info_rank0(
1054 "User callback registered: %s", type(callback).__name__,
1055 )
1057 # ------------------------------------------------------------------
1058 # Callback dispatch (explicit mode)
1059 # ------------------------------------------------------------------
1061 def _builtin_callbacks(self) -> list:
1062 """Return built-in callbacks in fixed dispatch order.
1064 Centralised so every dispatcher iterates the same list — adding a
1065 callback only needs an entry here plus a named field in
1066 ``_init_callbacks`` (no per-event copy/paste).
1067 """
1068 return [
1069 self.logging_callback,
1070 self.eval_callback,
1071 self.profiler_callback,
1072 self.wandb_callback,
1073 self.tensorboard_callback,
1074 self.progress_callback,
1075 self.checkpoint_callback,
1076 self.hf_export_callback,
1077 self.moe_monitor_callback,
1078 self.training_state_monitor_callback,
1079 self.gradient_health_callback,
1080 self.memory_monitor_callback,
1081 self.gc_callback,
1082 ]
1084 def _all_callbacks(self) -> list:
1085 """Built-in callbacks followed by user-registered ones."""
1086 return self._builtin_callbacks() + list(self.user_callbacks)
1088 def on_init_end(self):
1089 """Dispatch one-shot ``on_init_end`` after every ``_build_*`` ran.
1091 Fired by the subclass at the end of its own ``__init__`` (see
1092 ``LLMTrainer.__init__``); ``BaseTrainer.train()`` does NOT call it
1093 because BaseTrainer instances are sometimes wrapped (composition
1094 pattern) and the wrapper owns the init lifecycle.
1095 """
1096 for cb in self._all_callbacks():
1097 cb.on_init_end(self.state)
1099 def on_train_begin(self):
1100 """Dispatch on_train_begin to all callbacks."""
1101 # Memory monitor first so it captures the truly-initial peak.
1102 self.memory_monitor_callback.on_train_begin(self.state)
1103 self.moe_monitor_callback.on_train_begin(self.state)
1104 self.training_state_monitor_callback.on_train_begin(self.state)
1105 self.profiler_callback.on_train_begin(self.state)
1106 self.wandb_callback.on_train_begin(self.state)
1107 self.tensorboard_callback.on_train_begin(self.state)
1108 # Checkpoint runs after log writers are armed and before progress so
1109 # resumed ``global_step`` is reflected in the tqdm initial position.
1110 self.checkpoint_callback.on_train_begin(self.state)
1111 self.progress_callback.on_train_begin(self.state)
1112 for cb in self.user_callbacks:
1113 cb.on_train_begin(self.state)
1115 def on_train_end(self):
1116 """Dispatch on_train_end to all callbacks."""
1117 self.checkpoint_callback.on_train_end(self.state)
1118 self.hf_export_callback.on_train_end(self.state)
1119 self.progress_callback.on_train_end(self.state)
1120 self.training_state_monitor_callback.on_train_end(self.state)
1121 self.tensorboard_callback.on_train_end(self.state)
1122 self.wandb_callback.on_train_end(self.state)
1123 self.profiler_callback.on_train_end(self.state)
1124 for cb in self.user_callbacks:
1125 cb.on_train_end(self.state)
1127 def on_step_begin(self):
1128 """Dispatch on_step_begin to all callbacks."""
1129 self.logging_callback.on_step_begin(self.state)
1130 for cb in self.user_callbacks:
1131 cb.on_step_begin(self.state)
1133 def on_step_end(self, loss=None, grad_norm=None):
1134 """Dispatch on_step_end to all callbacks (built-ins + user)."""
1135 self.training_state_monitor_callback.on_step_end(
1136 self.state, loss=loss, grad_norm=grad_norm,
1137 )
1138 for cb in self._all_callbacks():
1139 if cb is self.training_state_monitor_callback:
1140 continue
1141 cb.on_step_end(self.state, loss=loss, grad_norm=grad_norm)
1143 def on_substep_end(self):
1144 """Dispatch on_substep_end (after each micro-batch forward/backward)."""
1145 self.moe_monitor_callback.on_substep_end(self.state)
1146 self.training_state_monitor_callback.on_substep_end(self.state)
1147 for cb in self.user_callbacks:
1148 cb.on_substep_end(self.state)
1150 def on_pre_optimizer_step(self, grad_norm=None):
1151 """Dispatch on_pre_optimizer_step (after grad clip, before optimizer.step)."""
1152 # Health check runs FIRST so a NaN aborts before the logger misleads.
1153 self.training_state_monitor_callback.on_pre_optimizer_step(
1154 self.state, grad_norm=grad_norm,
1155 )
1156 self.gradient_health_callback.on_pre_optimizer_step(
1157 self.state, grad_norm=grad_norm,
1158 )
1159 self.logging_callback.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
1160 self.wandb_callback.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
1161 self.tensorboard_callback.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
1162 for cb in self.user_callbacks:
1163 cb.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
1165 def on_epoch_begin(self):
1166 """Dispatch on_epoch_begin."""
1167 for cb in self._all_callbacks():
1168 cb.on_epoch_begin(self.state)
1170 def on_epoch_end(self):
1171 """Dispatch on_epoch_end."""
1172 for cb in self._all_callbacks():
1173 cb.on_epoch_end(self.state)
1175 # ------------------------------------------------------------------
1176 # Event fan-out (LoggingCallback / CheckpointCallback emit these)
1177 # ------------------------------------------------------------------
1179 def dispatch_log_event(self, metrics: dict) -> None:
1180 """Forward a metrics record to every callback's ``on_log``.
1182 ``LoggingCallback`` calls this so TensorBoard / W&B / external sinks
1183 log the SAME numbers — single source of truth, no duplicate work.
1184 """
1185 for cb in self._all_callbacks():
1186 cb.on_log(self.state, metrics=metrics)
1188 def dispatch_save_event(self, checkpoint_dir: str) -> None:
1189 """Forward a ckpt-save event to every callback's ``on_save``."""
1190 for cb in self._all_callbacks():
1191 cb.on_save(self.state, checkpoint_dir=checkpoint_dir)
1193 def dispatch_load_event(self, checkpoint_dir: str) -> None:
1194 """Forward a ckpt-load event to every callback's ``on_load``."""
1195 for cb in self._all_callbacks():
1196 cb.on_load(self.state, checkpoint_dir=checkpoint_dir)
1198 def dispatch_evaluate_event(self, metrics: dict = None) -> None:
1199 """Forward an eval-pass-complete event to every callback's ``on_evaluate``."""
1200 for cb in self._all_callbacks():
1201 cb.on_evaluate(self.state, metrics=metrics)
1203 # ------------------------------------------------------------------
1204 # Training core
1205 # ------------------------------------------------------------------
1207 def _move_value_to_device(self, value):
1208 """Move nested tensor-like values to this trainer's device."""
1209 if hasattr(value, "to"):
1210 return value.to(self.device, non_blocking=True)
1211 if isinstance(value, dict):
1212 return {k: self._move_value_to_device(v) for k, v in value.items()}
1213 if isinstance(value, list):
1214 return [self._move_value_to_device(v) for v in value]
1215 if isinstance(value, tuple):
1216 return tuple(self._move_value_to_device(v) for v in value)
1217 return value
1219 def _prepare_forward_batch(self, micro_batch):
1220 """Move a micro-batch to device and extract CP-shifted labels."""
1221 micro_batch = {
1222 key: self._move_value_to_device(value)
1223 for key, value in micro_batch.items()
1224 }
1225 labels_are_shifted = bool(micro_batch.pop("_hp_labels_are_shifted", False))
1226 shifted_labels = micro_batch.pop("labels", None) if labels_are_shifted else None
1227 if labels_are_shifted and shifted_labels is None:
1228 raise ValueError("CP-shifted loss marker is set but labels are missing.")
1229 return micro_batch, labels_are_shifted, shifted_labels
1231 def _compute_micro_loss(
1232 self,
1233 outputs,
1234 labels_are_shifted: bool,
1235 shifted_labels,
1236 micro_batch_tokens: int,
1237 ):
1238 """Return mean loss and summed loss for one micro-batch."""
1239 if not labels_are_shifted:
1240 loss = outputs["loss"] if isinstance(outputs, dict) else outputs.loss
1241 return loss, loss.detach() * max(micro_batch_tokens, 1)
1243 logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
1244 target_device = logits.device if hasattr(logits, "device") else self.device
1245 shifted_labels = shifted_labels.to(target_device, non_blocking=True)
1246 loss_sum = torch.nn.functional.cross_entropy(
1247 logits.float().view(-1, logits.size(-1)),
1248 shifted_labels.contiguous().view(-1),
1249 ignore_index=-100,
1250 reduction="sum",
1251 )
1252 return loss_sum / max(micro_batch_tokens, 1), loss_sum
1254 def _scale_loss_for_backward(
1255 self,
1256 loss,
1257 loss_sum,
1258 labels_are_shifted: bool,
1259 micro_batch_tokens: int,
1260 global_tokens: int,
1261 num_micro: int,
1262 ):
1263 """Scale one micro-batch loss according to trainer loss aggregation."""
1264 dp_size = self.parallel_dims.dp_size
1265 agg = self.args.train.optimizer.loss_aggregation
1266 cp_size = self._cp_size()
1267 cp_rank_average = agg == "rank_average" and cp_size > 1
1268 if agg == 'rank_average' and not cp_rank_average:
1269 scaled_loss = loss / num_micro if num_micro > 1 else loss
1270 rank_average_loss_scale_size = getattr(
1271 self.model,
1272 "hp_rank_average_loss_scale_size",
1273 1,
1274 )
1275 if rank_average_loss_scale_size != 1:
1276 scaled_loss = scaled_loss / rank_average_loss_scale_size
1277 return scaled_loss
1279 loss_scale_size = getattr(self.model, "hp_token_loss_scale_size", dp_size)
1280 if labels_are_shifted:
1281 scaled_loss = loss_sum / max(global_tokens, 1) * loss_scale_size
1282 else:
1283 scaled_loss = mean_global_loss(
1284 loss, micro_batch_tokens, global_tokens, loss_scale_size,
1285 )
1286 tp_loss_scale_size = getattr(
1287 self.model,
1288 "hp_loss_tp_scale_size",
1289 max(1, self._parallel_dim_size("tp")),
1290 )
1291 if tp_loss_scale_size != 1:
1292 scaled_loss = scaled_loss / tp_loss_scale_size
1293 ep_loss_scale_size = getattr(self.model, "hp_loss_ep_scale_size", 1)
1294 if ep_loss_scale_size != 1:
1295 scaled_loss = scaled_loss / ep_loss_scale_size
1296 return scaled_loss
1298 def forward_backward_step(
1299 self,
1300 micro_batch: Dict[str, Any],
1301 micro_batch_tokens: int,
1302 global_tokens: int,
1303 num_micro: int = 1,
1304 ):
1305 """Run forward + backward for one micro-batch.
1307 Uses global token normalisation: each micro-batch's
1308 loss is scaled by ``micro_tokens / global_tokens`` so that every token
1309 across all ranks and all micro-batches contributes equally to the
1310 gradient, regardless of DP size or grad_accum.
1312 Args:
1313 micro_batch: Dict of input tensors.
1314 micro_batch_tokens: Non-padding token count for this micro-batch.
1315 global_tokens: Total non-padding tokens across **all** ranks and
1316 **all** micro-batches (computed via all-reduce).
1318 Returns:
1319 Tuple of (raw_loss_scalar, micro_batch_tokens) for logging.
1320 """
1321 micro_batch, labels_are_shifted, shifted_labels = self._prepare_forward_batch(micro_batch)
1323 # Forward (with training context for activation offload)
1324 with self.model_fwd_context:
1325 outputs = self.model(**micro_batch, use_cache=False)
1326 loss, loss_sum = self._compute_micro_loss(
1327 outputs, labels_are_shifted, shifted_labels, micro_batch_tokens,
1328 )
1330 # TP scenario: loss may be Partial DTensor — reduce before backward
1331 if hasattr(loss, 'is_partial') and loss.is_partial():
1332 loss = loss.reduce_partial()
1334 # Keep raw loss value for logging before scaling
1335 raw_loss = loss.detach()
1337 scaled_loss = self._scale_loss_for_backward(
1338 loss,
1339 loss_sum,
1340 labels_are_shifted,
1341 micro_batch_tokens,
1342 global_tokens,
1343 num_micro,
1344 )
1346 # Backward (with training context)
1347 with self.model_bwd_context:
1348 scaled_loss.backward()
1350 return raw_loss, micro_batch_tokens
1352 def _shard_micro_batches_for_cp(self, micro_batches):
1353 """Slice each micro-batch's sequence onto this context-parallel rank.
1355 Under CP the model forward consumes only this rank's sequence slice (the
1356 Ulysses all-to-all / sequence-gather reconstruct the full sequence inside
1357 attention). The next-token shift is performed here on the **full**
1358 sequence before slicing so the cross-rank boundary target is preserved.
1359 The model remains HF-like: it receives explicit global ``position_ids``
1360 and no CP-only forward arguments. The trainer computes cross-entropy
1361 from the model logits for these pre-shifted local targets, and the
1362 per-rank token counts aggregate back to the single-card loss across the
1363 ``cp`` group (folded into the trainer's loss / FSDP reduction). No-op
1364 when ``cp<=1``.
1366 Args:
1367 micro_batches: List of per-micro-batch dicts from the data iterator.
1369 Returns:
1370 The CP-sharded micro-batch list (or the input unchanged when ``cp<=1``).
1371 """
1372 cp_size = self._cp_size()
1373 if cp_size <= 1:
1374 return micro_batches
1375 cp_rank = self.mesh["cp"].get_local_rank()
1376 sharded = []
1377 for micro_batch in micro_batches:
1378 input_ids = micro_batch["input_ids"]
1379 seq_len = input_ids.shape[1]
1380 if seq_len % cp_size != 0:
1381 raise ValueError(
1382 f"sequence length ({seq_len}) must be divisible by cp ({cp_size})."
1383 )
1384 shard = seq_len // cp_size
1385 start = cp_rank * shard
1386 seq_slice = slice(start, start + shard)
1387 local = dict(micro_batch)
1388 local["input_ids"] = input_ids[:, seq_slice].contiguous()
1389 position_ids = micro_batch.get("position_ids")
1390 if position_ids is not None:
1391 if position_ids.dim() == 2:
1392 local["position_ids"] = position_ids[:, seq_slice].contiguous()
1393 else:
1394 pos_slice = [slice(None)] * position_ids.dim()
1395 pos_slice[-1] = seq_slice
1396 local["position_ids"] = position_ids[tuple(pos_slice)].contiguous()
1397 else:
1398 has_multimodal_positions = any(
1399 micro_batch.get(name) is not None
1400 for name in (
1401 "pixel_values", "image_grid_thw", "pixel_values_videos",
1402 "video_grid_thw", "mm_token_type_ids",
1403 )
1404 )
1405 if not has_multimodal_positions:
1406 local["position_ids"] = torch.arange(
1407 start, start + shard, device=input_ids.device, dtype=torch.long,
1408 ).view(1, -1).expand(input_ids.shape[0], -1)
1409 mm_token_type_ids = micro_batch.get("mm_token_type_ids")
1410 if mm_token_type_ids is not None:
1411 mm_slice = [slice(None)] * mm_token_type_ids.dim()
1412 mm_slice[-1] = seq_slice
1413 local["mm_token_type_ids"] = mm_token_type_ids[
1414 tuple(mm_slice)
1415 ].contiguous()
1416 labels = micro_batch.get("labels")
1417 if labels is not None:
1418 shifted = torch.nn.functional.pad(labels, (0, 1), value=-100)[..., 1:]
1419 local["labels"] = shifted[:, seq_slice].contiguous()
1420 local["_hp_labels_are_shifted"] = True
1421 attn = micro_batch.get("attention_mask")
1422 if attn is not None and hasattr(attn, "dim") and attn.dim() == 2:
1423 local["attention_mask"] = attn[:, seq_slice].contiguous()
1424 sharded.append(local)
1425 return sharded
1427 def _collect_global_tokens(self, token_counts):
1428 """Count valid loss tokens and all-reduce across the data-parallel group."""
1429 local_tokens = sum(token_counts) or 1
1430 global_tokens = local_tokens
1431 if platform.get_world_size() > 1 and self._dp_group_info.group is not None:
1432 token_tensor = platform.full((1,), local_tokens).to(self.device)
1433 platform.all_reduce(token_tensor, self._dp_group_info)
1434 global_tokens = max(int(token_tensor.item()), 1)
1435 self._last_global_tokens = global_tokens
1436 return local_tokens, global_tokens
1438 def _run_micro_batches(self, micro_batches, token_counts, global_tokens):
1439 """Run forward/backward over accumulated micro-batches."""
1440 num_micro = len(micro_batches)
1441 total_loss_sum = 0.0
1442 total_loss_arith_sum = 0.0
1443 total_tokens_local = 0
1444 for index, micro_batch in enumerate(micro_batches):
1445 is_last = index == num_micro - 1
1446 if isinstance(self.model, HSDPModule):
1447 self.model.set_requires_gradient_sync(is_last)
1448 self.model.set_is_last_backward(is_last)
1449 self._maybe_toggle_reshard(index, num_micro)
1451 raw_loss, micro_tokens = self.forward_backward_step(
1452 micro_batch,
1453 token_counts[index],
1454 global_tokens,
1455 num_micro=num_micro,
1456 )
1457 loss_value = raw_loss.item()
1458 total_loss_sum += loss_value * micro_tokens
1459 total_loss_arith_sum += loss_value
1460 total_tokens_local += micro_tokens
1461 self.state.substep_info = {
1462 "raw_loss": loss_value,
1463 "micro_tokens": micro_tokens,
1464 }
1465 self.on_substep_end()
1466 return total_loss_sum, total_loss_arith_sum, total_tokens_local
1468 def _run_post_fsdp_grad_reduce(self) -> None:
1469 """Run an optional model-provided reducer after FSDP gradients drain."""
1470 post_fsdp_grad_reduce = getattr(self.model, "hp_post_fsdp_grad_reduce", None)
1471 if post_fsdp_grad_reduce is not None:
1472 post_fsdp_grad_reduce()
1474 def _non_pp_clip_grad_norm(self, max_grad_norm: float):
1475 """Clip non-pipeline gradients using the configured clipping function."""
1476 clip_fn = self.spec.clip_grad_fn or clip_grad_norm_
1477 return clip_fn(self.model.parameters(), max_grad_norm)
1479 def _optimizer_step_after_backward(self, clip_fn):
1480 """Clip gradients if enabled, run optimizer/scheduler, and clear grads."""
1481 max_grad_norm = float(self.args.train.optimizer.max_grad_norm)
1482 grad_norm = clip_fn(max_grad_norm) if max_grad_norm > 0.0 else None
1483 grad_norm_value = None if grad_norm is None else grad_norm.item()
1484 self.on_pre_optimizer_step(grad_norm=grad_norm_value)
1486 with SkipDTensorDispatch():
1487 self.optimizer.step()
1488 if self.lr_scheduler is not None:
1489 self.lr_scheduler.step()
1490 self.optimizer.zero_grad()
1491 return grad_norm_value
1493 def _aggregate_non_pp_loss(
1494 self,
1495 total_loss_sum: float,
1496 total_loss_arith_sum: float,
1497 total_tokens_local: int,
1498 global_tokens: int,
1499 num_micro: int,
1500 ) -> float:
1501 """Aggregate the reported non-pipeline loss across DP ranks."""
1502 agg = self.args.train.optimizer.loss_aggregation
1503 cp_size = self._cp_size()
1504 if agg == "token_weighted" or (agg == "rank_average" and cp_size > 1):
1505 if platform.get_world_size() > 1 and self._dp_group_info.group is not None:
1506 loss_tensor = platform.full((1,), total_loss_sum).to(self.device)
1507 platform.all_reduce(loss_tensor, self._dp_group_info)
1508 return loss_tensor.item() / max(global_tokens, 1)
1509 return total_loss_sum / max(total_tokens_local, 1)
1511 local_mean = total_loss_arith_sum / max(num_micro, 1)
1512 dp_size = self._dp_group_info.rank_size
1513 if dp_size <= 1:
1514 return local_mean
1515 loss_tensor = platform.full((1,), local_mean).to(self.device)
1516 platform.all_reduce(loss_tensor, self._dp_group_info)
1517 return loss_tensor.item() / dp_size
1519 def _average_model_parallel_metric(self, avg_loss: float) -> float:
1520 """Average replicated loss metrics over model-parallel EP when needed."""
1521 tp_size = self._parallel_dim_size("tp")
1522 ep_size = self._parallel_dim_size("ep")
1523 if tp_size > 1 and ep_size > 1:
1524 return avg_loss
1525 if ep_size <= 1:
1526 return avg_loss
1527 try:
1528 ep_group = self.mesh.get_group("ep")
1529 except (KeyError, ValueError):
1530 return avg_loss
1531 metric = platform.full((1,), avg_loss).to(self.device)
1532 ep_group_info = GroupInfo(
1533 group_name="trainer_ep_metric",
1534 group=ep_group,
1535 rank_size=ep_size,
1536 )
1537 platform.all_reduce(metric, ep_group_info)
1538 return metric.item() / ep_size
1540 def train_step(self, data_iterator):
1541 """Execute one training step with gradient accumulation.
1543 Consistent across different DP configurations by:
1544 1. All-reducing global token count before loss scaling ()
1545 2. Syncing gradients only on the last micro-batch ()
1546 3. All-reducing loss weighted by token count for reporting
1548 Args:
1549 data_iterator: Iterator yielding lists of micro-batch dicts.
1550 """
1551 if self.pp_enabled:
1552 return self._pp_train_step(data_iterator)
1553 micro_batches = next(data_iterator)
1554 prepare_batch_fn = getattr(self.spec, "prepare_batch_fn", None)
1555 if prepare_batch_fn is not None:
1556 micro_batches = [
1557 prepare_batch_fn(batch, self.model)
1558 for batch in micro_batches
1559 ]
1560 micro_batches = self._shard_micro_batches_for_cp(micro_batches)
1561 self.state.global_step += 1
1562 num_micro = len(micro_batches)
1564 token_counts = [count_loss_token(mb) for mb in micro_batches]
1565 _, global_tokens = self._collect_global_tokens(token_counts)
1566 total_loss_sum, total_loss_arith_sum, total_tokens_local = self._run_micro_batches(
1567 micro_batches,
1568 token_counts,
1569 global_tokens,
1570 )
1572 # Wait for async gradient reduce
1573 #
1574 hsdp_sync_stream()
1575 self._run_post_fsdp_grad_reduce()
1576 grad_norm_value = self._optimizer_step_after_backward(self._non_pp_clip_grad_norm)
1577 avg_loss = self._aggregate_non_pp_loss(
1578 total_loss_sum,
1579 total_loss_arith_sum,
1580 total_tokens_local,
1581 global_tokens,
1582 num_micro,
1583 )
1584 avg_loss = self._average_model_parallel_metric(avg_loss)
1586 return {"loss": avg_loss, "grad_norm": grad_norm_value}
1588 @staticmethod
1589 def _pp_concat_micro_batches(micro_batches):
1590 """Concatenate grad-accum micro-batches into one global batch (dim 0).
1592 Under PP the schedule owns micro-batching, so the trainer rebuilds the
1593 global batch from the grad-accum group and lets ``ScheduleGPipe``
1594 re-split it into ``pp_micro_batch_num`` chunks.
1596 The pipeline runs a single fused ``sum``-CE backward over the whole
1597 batch, which reproduces the trainer's ``token_weighted`` single-card
1598 gradient **only when every micro-batch shares the same sequence length**
1599 (then ``sum-CE / valid_tokens`` is the common token-mean). Micro-batches
1600 of differing shape are therefore rejected with a clear error — pad to a
1601 fixed ``max_seq_len`` so the grad-accum group is uniform, or size the
1602 batch so ``grad_accum == 1``. Non-tensor values are taken from the first
1603 micro-batch.
1604 """
1605 if len(micro_batches) == 1:
1606 return dict(micro_batches[0])
1607 merged = {}
1608 for key in micro_batches[0].keys():
1609 values = [mb[key] for mb in micro_batches]
1610 first = values[0]
1611 if not hasattr(first, "dim"):
1612 merged[key] = first
1613 continue
1614 if any(value.shape[1:] != first.shape[1:] for value in values):
1615 raise NotImplementedError(
1616 f"PP gradient accumulation requires uniform-shape "
1617 f"micro-batches; '{key}' varies across the group (shapes "
1618 f"{[tuple(value.shape) for value in values]}). Pad to a fixed "
1619 f"max_seq_len, or size the batch so grad_accum == 1."
1620 )
1621 merged[key] = torch.cat(values, dim=0)
1622 return merged
1624 def _pp_clip_grad_norm(self, max_grad_norm: float):
1625 """Clip gradients by the **global** norm across all pipeline stages.
1627 Each stage holds a disjoint parameter slab, so the single-card total
1628 norm is recovered by summing the per-stage squared norms and all-reducing
1629 over the pipeline group. The shared coefficient is then applied on every
1630 stage — essential for the tied embed / lm_head, whose stage-0 and
1631 last-stage copies must receive the *same* scaling to stay bit-identical
1632 after the optimizer step (a per-stage coefficient would desync them).
1634 The tied copy is counted once: the last stage skips its ``lm_head.weight``
1635 duplicate from the norm sum (it equals stage 0's ``embed_tokens.weight``)
1636 but is still scaled, so the global norm matches the single-card norm.
1638 Args:
1639 max_grad_norm: Clip threshold; the effective coefficient is
1640 ``min(1, max_grad_norm / total_norm)``.
1642 Returns:
1643 The global gradient norm (a scalar tensor) for logging.
1644 """
1645 params = [p for p in self.model.parameters() if p.grad is not None]
1646 skip = None
1647 if self._pp_tie_embeddings and self.pp_has_last_stage:
1648 # The last global stage's submodule owns the tied ``lm_head``. Under
1649 # VPP ``self.model`` is a ModuleList of this rank's chunks, only one
1650 # of which (the last stage) carries ``lm_head`` — find it there.
1651 head_owner = self.model
1652 if isinstance(head_owner, torch.nn.ModuleList):
1653 head_owner = next(
1654 (s for s in head_owner if hasattr(s, "lm_head")), None)
1655 if head_owner is not None and hasattr(head_owner, "lm_head"):
1656 skip = head_owner.lm_head.weight
1657 local_sq = torch.zeros((), device=self.device, dtype=torch.float32)
1658 for param in params:
1659 if param is skip:
1660 continue
1661 grad = param.grad.detach()
1662 # Under PP+FSDP the grad is a sharded DTensor; reduce on the local
1663 # shard so the cross-stage all-reduce stays a plain-tensor collective.
1664 if hasattr(grad, "to_local"):
1665 grad = grad.to_local()
1666 local_sq = local_sq + grad.float().pow(2).sum()
1667 platform.all_reduce(local_sq, self._pp_group_info)
1668 # Under PP+FSDP the grads are dp-sharded, so also sum the per-dp-shard
1669 # squared norms across the dp group to get the true global grad norm.
1670 if getattr(self, "_pp_stage_fsdp_sharded", False):
1671 platform.all_reduce(local_sq, self._dp_group_info)
1672 total_norm = local_sq.sqrt()
1673 clip_coef = (max_grad_norm / (total_norm + 1e-6)).clamp(max=1.0)
1674 for param in params:
1675 param.grad.mul_(clip_coef.to(param.grad.dtype))
1676 return total_norm
1678 def _pp_load_first_stage_batch(self, data_iterator):
1679 """Load and prepare the global PP batch on the first stage only."""
1680 batch = None
1681 targets = None
1682 stop = 0
1683 if not self.pp_has_first_stage:
1684 return batch, targets, stop
1685 try:
1686 micro_batches = next(data_iterator)
1687 batch = self._pp_concat_micro_batches(micro_batches)
1688 batch = {
1689 key: (value.to(self.device, non_blocking=True) if hasattr(value, "to") else value)
1690 for key, value in batch.items()
1691 }
1692 if batch["input_ids"].shape[0] % self.pp_micro_batch_num != 0:
1693 stop = 1
1694 else:
1695 labels = batch["labels"]
1696 targets = torch.nn.functional.pad(labels, (0, 1), value=-100)[..., 1:].to(torch.int64)
1697 except StopIteration:
1698 stop = 1
1699 return batch, targets, stop
1701 def _pp_broadcast_control(self, batch, targets, stop: int):
1702 """Broadcast stop/shape metadata across the pipeline group."""
1703 ctrl = platform.full((4,), 0, dtype=torch.int64).to(self.device)
1704 if stop:
1705 ctrl[0] = 1
1706 elif self.pp_has_first_stage:
1707 ctrl[1] = int(targets.shape[0])
1708 ctrl[2] = int(targets.shape[1])
1709 ctrl[3] = 1 if batch.get("attention_mask") is not None else 0
1710 platform.broadcast(ctrl, self._pp_src_rank, self._pp_group_info.group)
1711 return ctrl.tolist()
1713 def _pp_broadcast_2d_int64(self, src_tensor, rows: int, seq: int):
1714 """Broadcast one 2-D int64 tensor from the first pipeline stage."""
1715 tensor = (
1716 src_tensor.to(torch.int64).contiguous()
1717 if self.pp_has_first_stage
1718 else platform.full((rows, seq), 0, dtype=torch.int64).to(self.device)
1719 )
1720 platform.broadcast(tensor, self._pp_src_rank, self._pp_group_info.group)
1721 return tensor
1723 def _pp_prepare_broadcast_inputs(self, batch, targets, stop: int):
1724 """Broadcast targets and optional all-stage masks for one PP step."""
1725 stop, rows, seq, has_attn = self._pp_broadcast_control(batch, targets, stop)
1726 if stop:
1727 raise StopIteration
1729 targets = self._pp_broadcast_2d_int64(targets, rows, seq)
1730 attention_mask = None
1731 if has_attn:
1732 source_mask = batch["attention_mask"] if self.pp_has_first_stage else None
1733 attention_mask = self._pp_broadcast_2d_int64(source_mask, rows, seq)
1734 return targets, attention_mask, has_attn
1736 def _pp_count_valid_tokens(self, targets) -> int:
1737 """Count valid shifted targets and sum across DP for PP+FSDP."""
1738 n_valid = max(int((targets != -100).sum().item()), 1)
1739 if getattr(self, "_pp_fsdp_composed", False):
1740 token_tensor = platform.full((1,), n_valid).to(self.device)
1741 platform.all_reduce(token_tensor, self._dp_group_info)
1742 n_valid = max(int(token_tensor.item()), 1)
1743 self._last_global_tokens = n_valid
1744 return n_valid
1746 def _pp_validate_rank_average_targets(self, targets) -> None:
1747 """Validate the PP token-mean path also represents rank-average loss."""
1748 agg = self.args.train.optimizer.loss_aggregation
1749 if agg != "rank_average":
1750 return
1751 row_tokens = (targets != -100).sum(dim=1)
1752 if row_tokens.numel() <= 1:
1753 return
1754 if int(row_tokens.min().item()) == int(row_tokens.max().item()):
1755 return
1756 raise NotImplementedError(
1757 "Trainer PP with loss_aggregation='rank_average' requires uniform "
1758 "valid-token counts per row so the fused token-mean loss matches "
1759 "the single-card rank-average gradient."
1760 )
1762 def _pp_normalize_grads(self, n_valid: int) -> None:
1763 """Normalize fully reduced pipeline gradients to the global token mean.
1765 Core pipeline schedules retain unit backward sensitivity for standalone
1766 callers. After PP/FSDP/shared/TP/EP/DP reductions, multiplying the final
1767 averaged gradients by ``dp_size / (n_valid * tp_loss_repeats)`` yields
1768 the same global token mean before clipping and the optimizer step. The
1769 TP divisor removes duplicate backward sensitivity when the last stage
1770 materializes a replicated loss as a local tensor.
1771 """
1772 dp_size = max(int(self.parallel_dims.dp_size), 1)
1773 denominator = max(n_valid * self._pp_tp_loss_repeats, 1)
1774 grad_scale = dp_size / denominator
1775 for param in self.model.parameters():
1776 if not param.requires_grad:
1777 continue
1778 grad = getattr(param, "main_grad", None)
1779 if grad is None:
1780 grad = param.grad
1781 if grad is None:
1782 continue
1783 local_grad = grad.to_local() if isinstance(grad, DTensor) else grad
1784 local_grad.mul_(grad_scale)
1786 def _pp_run_schedule(self, batch, targets, attention_mask, has_attn):
1787 """Run the configured PP schedule with the broadcast inputs."""
1788 run_kwargs = {"targets": targets}
1789 kwargs_batch_dim = getattr(self.pp_schedule, "_kwargs_batch_dim", {}) or {}
1790 if self.pp_has_first_stage:
1791 for key in kwargs_batch_dim:
1792 if key != "targets" and key in batch:
1793 run_kwargs[key] = batch[key]
1794 return self.pp_schedule.run(batch["input_ids"], **run_kwargs)
1795 if has_attn and "attention_mask" in kwargs_batch_dim:
1796 run_kwargs["attention_mask"] = attention_mask
1797 return self.pp_schedule.run(**run_kwargs)
1799 def _pp_post_schedule_grad_reduce(self) -> None:
1800 """Run optional post-FSDP reducers on local pipeline stage modules."""
1801 stage_modules = list(self.model) if isinstance(self.model, torch.nn.ModuleList) else [self.model]
1802 for stage_module in stage_modules:
1803 stage_tp_reduce = getattr(stage_module, "hp_post_fsdp_grad_reduce", None)
1804 if stage_tp_reduce is not None:
1805 stage_tp_reduce()
1807 def _pp_average_plain_dp_grads(self) -> None:
1808 """Average plain replicated grads for PP+DP without per-stage FSDP shards."""
1809 if not getattr(self, "_pp_fsdp_composed", False):
1810 return
1811 dp_size = max(int(self.parallel_dims.dp_size), 1)
1812 if dp_size <= 1 or getattr(self, "_pp_stage_fsdp_sharded", False):
1813 return
1814 for param in self.model.parameters():
1815 if param.grad is not None:
1816 platform.all_reduce(param.grad, self._dp_group_info)
1817 param.grad.div_(dp_size)
1819 def _pp_reduce_reported_loss(self, outputs, n_valid: int) -> float:
1820 """Reduce last-stage sum-CE into a reported token-mean PP loss."""
1821 local_sum_ce = 0.0
1822 if self.pp_has_last_stage:
1823 local_sum_ce = sum(out.detach().float() for out in outputs).item()
1824 sum_ce_t = platform.full((1,), local_sum_ce).to(self.device)
1825 if getattr(self, "_pp_fsdp_composed", False):
1826 platform.all_reduce(sum_ce_t, self._dp_group_info)
1827 loss_t = sum_ce_t / n_valid
1828 platform.all_reduce(loss_t, self._pp_group_info)
1829 return loss_t.item()
1831 def _pp_train_step(self, data_iterator):
1832 """Pipeline-parallel training step (``pp > 1``).
1834 Only the first stage reads the dataloader; the last stage's ``targets``
1835 and the all-stage ``attention_mask`` are broadcast across the pipeline
1836 group so non-first stages never load (and, for VL, never decode) the
1837 identical batch. Heavy vision inputs stay on stage 0.
1839 ``ScheduleGPipe`` owns micro-batching and the forward/backward, so the
1840 trainer feeds it the **full** global batch (the grad-accum micro-batches
1841 concatenated). Only the last stage produces the per-micro-batch sum-CE;
1842 it is normalised to mean-CE and all-reduced across the pipeline group so
1843 every rank — including the rank-0 logger, which is the *first* stage —
1844 reports the same loss matching the single-card token-mean baseline.
1845 Gradient clipping uses the **global** cross-stage norm
1846 (:meth:`_pp_clip_grad_norm`) so every stage scales by the same
1847 coefficient — required so the tied embed / lm_head copies stay in sync.
1848 """
1849 batch, targets, stop = self._pp_load_first_stage_batch(data_iterator)
1850 targets, attention_mask, has_attn = self._pp_prepare_broadcast_inputs(batch, targets, stop)
1851 self.state.global_step += 1
1852 self._pp_validate_rank_average_targets(targets)
1853 n_valid = self._pp_count_valid_tokens(targets)
1854 outputs = self._pp_run_schedule(batch, targets, attention_mask, has_attn)
1855 self._pp_post_schedule_grad_reduce()
1856 self._pp_average_plain_dp_grads()
1857 self._pp_normalize_grads(n_valid)
1858 grad_norm_value = self._optimizer_step_after_backward(self._pp_clip_grad_norm)
1859 return {"loss": self._pp_reduce_reported_loss(outputs, n_valid), "grad_norm": grad_norm_value}
1861 def train(self):
1862 """Main training loop: epoch → step → micro-batch.
1864 Dispatches callbacks at each lifecycle point (explicit mode).
1865 on_train_begin is called first — CheckpointCallback uses it to restore
1866 state.global_step from a saved checkpoint, so the loop below will
1867 correctly skip already-completed steps.
1868 """
1869 logger.info_rank0(
1870 "Training starts: max_steps=%d, epochs=%d",
1871 self.state.max_steps,
1872 self.args.train.num_train_epochs,
1873 )
1874 # on_train_begin runs checkpoint resume — state.global_step may be
1875 # updated to the resumed step before the loop starts.
1876 self.on_train_begin()
1877 num_epochs = self.args.train.num_train_epochs
1879 if self.state.global_step > 0:
1880 logger.info_rank0(
1881 "Resuming training from step %d", self.state.global_step,
1882 )
1884 for epoch in range(num_epochs):
1885 if self.state.global_step >= self.state.max_steps:
1886 break
1887 self.state.epoch = epoch
1888 if hasattr(self, 'sampler'):
1889 self.sampler.set_epoch(epoch)
1890 self.on_epoch_begin()
1892 # Build micro-batch iterator from the stateful dataloader.
1893 # StatefulDataLoader tracks iterator position internally,
1894 # so after resume it skips already-consumed batches.
1895 data_iterator = self._make_micro_batch_iterator()
1897 # Drive the loop on the live ``global_step`` so total training
1898 # never exceeds ``max_steps`` regardless of ``num_train_epochs``
1899 # or resume offset.
1900 while self.state.global_step < self.state.max_steps:
1901 self.on_step_begin()
1902 try:
1903 metrics = self.train_step(data_iterator)
1904 except StopIteration:
1905 logger.info_rank0("Epoch %d: dataloader exhausted", epoch)
1906 break
1908 self.on_step_end(
1909 loss=metrics["loss"],
1910 grad_norm=metrics["grad_norm"],
1911 )
1913 self.on_epoch_end()
1915 self.on_train_end()
1916 destroy_process_group()
1917 logger.info_rank0("Training completed")
1919 # ------------------------------------------------------------------
1920 # Helpers
1921 # ------------------------------------------------------------------
1923 def _make_micro_batch_iterator(self):
1924 """Yield lists of micro-batches from the stateful dataloader.
1926 Groups ``self._grad_accum`` consecutive batches into a list for
1927 gradient accumulation. The underlying ``StatefulDataLoader`` tracks
1928 iteration position, so checkpoint/resume skips consumed batches.
1929 """
1930 batch_buffer = []
1931 for batch in self.train_dataloader:
1932 batch_buffer.append(batch)
1933 if len(batch_buffer) >= self._grad_accum:
1934 yield batch_buffer
1935 batch_buffer = []
1936 if batch_buffer:
1937 yield batch_buffer
1939 def _get_layers(self) -> list:
1940 """Return the repeating layers for FSDP/AC wrapping.
1942 Default: ``model.layers`` when the model exposes decoder layers.
1943 Override in subclass for models with different structure.
1944 """
1945 if hasattr(self.model, 'layers'):
1946 return list(self.model.layers)
1947 raise ValueError(
1948 f"Model {type(self.model).__name__} has no .layers attribute. "
1949 f"Either add self.layers to the model, or override _get_layers() "
1950 f"in the Trainer subclass."
1951 )
1953 def _get_combined_dp_group(self):
1954 """Return the combined data-parallel ProcessGroup for trainer all-reduce.
1956 Prefers the ``"loss"`` flatten alias registered by
1957 ``ParallelDims.build_mesh`` (folds CP into the DP group when CP is
1958 active so token-count denominators include CP-sharded contributions).
1959 Falls back to ``"dp"``, then to the legacy ``dp_shard`` /
1960 ``dp_replicate`` axes for callers that built a custom mesh.
1961 """
1962 for name in ("loss", "dp", "dp_shard", "dp_replicate"):
1963 try:
1964 return self.mesh.get_group(name)
1965 except (KeyError, ValueError):
1966 continue
1967 # No data-parallel axis: pure TP still needs the 1-D group because its
1968 # SequenceParallel ranks hold different token shards. Pure EP peers see
1969 # the same tokens and must not be folded into the token/loss denominator.
1970 if self.mesh.mesh_dim_names == ("ep",):
1971 return None
1972 # Other 1-D meshes (pure TP; pure CP normally has a ``loss`` alias)
1973 # return their own group. Multi-dim meshes with no DP/loss axis return
1974 # ``None``.
1975 try:
1976 return self.mesh.get_group()
1977 except (ValueError, RuntimeError):
1978 return None
1980 def _build_fsdp_kwargs(self) -> dict:
1981 """Build kwargs for ``fully_shard`` calls (dense parameters).
1983 For expert parameters when EP > 1, use ``_build_expert_fsdp_kwargs``.
1984 """
1985 for name in ("dp_shard", "dp", "dp_replicate"):
1986 try:
1987 dp_mesh = self.mesh[name]
1988 break
1989 except (KeyError, TypeError):
1990 continue
1991 else:
1992 dp_mesh = self.mesh
1993 kwargs = {"mesh": dp_mesh}
1995 reshard = self.args.train.accelerator.reshard_after_forward
1996 kwargs["reshard_after_forward"] = reshard
1998 return kwargs
2000 def _build_expert_fsdp_kwargs(self) -> dict:
2001 """Build kwargs for ``fully_shard`` calls on expert parameters.
2003 When EP > 1, expert parameters are sharded across the EP group
2004 with a separate mesh dimension. Falls back to dense FSDP kwargs
2005 if EP is not enabled.
2006 """
2007 if not self.parallel_dims.ep_enabled:
2008 return self._build_fsdp_kwargs()
2010 try:
2011 ep_mesh = self.mesh["ep"]
2012 except (KeyError, TypeError):
2013 logger.warning("EP=%d but no 'ep' dimension in mesh, falling back to dp mesh",
2014 self.parallel_dims.ep)
2015 return self._build_fsdp_kwargs()
2017 kwargs = {"mesh": ep_mesh}
2018 reshard = self.args.train.accelerator.reshard_after_forward
2019 kwargs["reshard_after_forward"] = reshard
2020 return kwargs
2022 def _materialize_and_init_shards(self) -> None:
2023 """Materialize meta-device parameters/buffers to real device in-place.
2025 After ``fully_shard`` on a meta-device model, each rank's parameters
2026 are meta DTensor shards **and FSDP2 holds internal views into those
2027 meta storages** (flat_param / unsharded buffer). Replacing the
2028 ``DTensor._local_tensor`` attribute leaves FSDP's internal views
2029 pointing at the old meta storage, so the first forward's all-gather
2030 still hits meta → ``c10d::_allgather_base_`` raises.
2032 PyTorch's ``nn.Module.to_empty(device=...)`` is the FSDP2-safe path:
2033 it walks every parameter/buffer (including DTensor shards) and
2034 **allocates real device storage in-place via ``torch.empty_like``**,
2035 preserving every existing view. After ``to_empty``, storage is
2036 uninitialised — we init on the local shard with kaiming_uniform for
2037 weights, zero for biases / 1-D / buffers.
2039 This is the meta-init path used after ``fully_shard`` has installed
2040 FSDP views.
2041 """
2042 device_type = platform.device_type()
2043 # Step 1: meta → real storage, in-place (FSDP-views preserved).
2044 self.model.to_empty(device=device_type)
2045 self._materialize_replicate_params(device_type)
2046 # Step 2: init the local shard of every param (and zero every buffer).
2047 param_count = self._init_local_shards()
2048 # Re-derive buffers wiped by ``to_empty`` (e.g. ``inv_freq``);
2049 # without this RoPE silently returns identity rotation.
2050 for module in self.model.modules():
2051 if hasattr(module, "reset_inv_freq"):
2052 module.reset_inv_freq()
2053 # Re-tie weights — ``to_empty`` gives every nn.Parameter fresh
2054 # storage so ``__init__``-time ties are broken. Must happen before
2055 # ``lazy_init`` re-wraps params as DTensor (non-leaf), which would
2056 # cause ``register_parameter`` to reject the assignment. Skipped under
2057 # PP: the tied embed / lm_head live on different stages, kept consistent
2058 # by the pipeline ``SharedParameterInfo`` (init broadcast + grad
2059 # all-reduce); a model-level tie would alias them into one object and
2060 # orphan the captured shared parameter (its grad would stay ``None``).
2061 if hasattr(self.model, "tie_weights") and int(self.parallel_dims.pp) <= 1:
2062 self.model.tie_weights()
2063 # ``to_empty`` strips DTensor; ``lazy_init`` re-wraps shards before
2064 # ``_load_weights`` / optimizer step see the params (the forward
2065 # pre-hook does the same later, but the loader needs DTensor first).
2066 reset_count = self._lazy_init_hsdp_modules()
2067 logger.info_rank0(
2068 "Meta → real on %s: to_empty + kaiming/zero init on %d params; "
2069 "FSDP lazy_init re-wrapped %d modules back to DTensor",
2070 device_type, param_count, reset_count,
2071 )
2073 def _iter_hsdp_states(self):
2074 """Yield the HSDP state attached to every HSDP-wrapped submodule."""
2075 seen = set()
2076 roots = [self.model, *getattr(self, "_pp_stage_modules", [])]
2077 for root in roots:
2078 if root is None:
2079 continue
2080 for module in root.modules():
2081 if not isinstance(module, HSDPModule):
2082 continue
2083 scheduler = getattr(module, 'hsdp_scheduler', None)
2084 state = getattr(scheduler, 'hsdp_state', None) if scheduler else None
2085 if state is None or id(state) in seen:
2086 continue
2087 seen.add(id(state))
2088 yield state
2090 def _materialize_replicate_params(self, device_type: str) -> None:
2091 """Materialize meta ``_local_tensor`` storage that ``to_empty`` cannot reach.
2093 Walks ``replicate_params`` (explicit no-shard buckets, e.g. ``(1, H)``
2094 shapes) and, for single-card FSDP, ``hsdp_params`` — the flat-buffer
2095 rebase in ``_init_flat_param_buffer`` is skipped at
2096 ``shard_world_size == 1``, leaving those params on meta and tripping
2097 ``_validate_no_meta_params`` in ``lazy_init``. The two buckets are
2098 disjoint by construction (see ``state.py`` ``_init_hsdp_params``).
2099 """
2100 for state in self._iter_hsdp_states():
2101 buckets = (
2102 getattr(state, 'replicate_params', []) or [],
2103 getattr(state, 'hsdp_params', []) or [],
2104 )
2105 for bucket in buckets:
2106 for hsdp_param in bucket:
2107 local = getattr(hsdp_param.sharded_param, "_local_tensor", None)
2108 if local is not None and local.is_meta:
2109 new_local = torch.empty_like(local, device=device_type)
2110 hsdp_param.sharded_param._local_tensor = new_local # pylint: disable=W0212
2111 hsdp_param._sharded_param_data = new_local.view(-1) # pylint: disable=W0212
2113 def _init_local_shards(self) -> int:
2114 """Init local shard of every param (kaiming for >=2D, zero else); zero buffers."""
2115 param_count = 0
2116 with torch.no_grad():
2117 for _, param in self.model.named_parameters():
2118 local = param._local_tensor if hasattr(param, '_local_tensor') else param # pylint: disable=W0212
2119 if local.is_meta:
2120 continue
2121 if local.dim() >= 2:
2122 torch.nn.init.kaiming_uniform_(local)
2123 else:
2124 torch.nn.init.zeros_(local)
2125 param_count += 1
2126 for _, buf in self.model.named_buffers():
2127 if buf is not None:
2128 buf.zero_()
2129 return param_count
2131 def _lazy_init_hsdp_modules(self) -> int:
2132 """Re-wrap HSDP shards into DTensor so loader / optimizer see them."""
2133 reset_count = 0
2134 for state in self._iter_hsdp_states():
2135 if hasattr(state, 'lazy_init'):
2136 state.lazy_init()
2137 reset_count += 1
2138 return reset_count
2140 def _load_weights(self, weights_path: str) -> None:
2141 """Load pre-trained weights from ``weights_path`` into the (possibly sharded) model.
2143 Uses hyper's distributed checkpoint ``load`` API so that each rank only
2144 reads the shard it owns. Falls back to a plain ``torch.load`` + partial
2145 ``load_state_dict`` for single-file checkpoints (e.g. safetensors).
2147 Args:
2148 weights_path: Path to a directory containing a distributed checkpoint,
2149 or a single ``.pt`` / ``.bin`` file.
2150 """
2151 logger.info_rank0("Loading weights from %s", weights_path)
2152 try:
2153 if os.path.isdir(weights_path):
2154 hf_index = os.path.join(weights_path, "model.safetensors.index.json")
2155 # Delegate model-specific renaming / expert-splitting to
2156 # the per-spec ``state_dict_adapter``.
2157 adapter_cls = getattr(self.spec, "state_dict_adapter", None)
2158 if os.path.isfile(hf_index) and adapter_cls is not None:
2159 self._load_hf_safetensors(weights_path, adapter_cls)
2160 else:
2161 self._load_hyper_dcp(weights_path)
2162 else:
2163 self._load_single_file(weights_path)
2164 logger.info_rank0("Weights loaded from %s", weights_path)
2165 except Exception as exc:
2166 raise RuntimeError(
2167 f"Failed to load weights from {weights_path}: {exc}. "
2168 "weights_path was provided so silent random-init fallback is unsafe — "
2169 "uniform-logits loss would corrupt downstream training metrics."
2170 ) from exc
2172 def _load_validated_state_dict(self, valid_sd: Dict[str, Any]) -> None:
2173 """Copy a validated plain-tensor state_dict into ``self.model``.
2175 Routes by model shape:
2177 * ``HSDPModule`` root (non-PP FSDP) — delegate to its shard-aware
2178 ``load_state_dict``, which distributes plain tensors onto local shards.
2179 * plain root with no DTensor params (no FSDP, or PP alone) — use the
2180 default ``load_state_dict`` (plain ``copy_``).
2181 * plain root that *holds* DTensor params (pipeline parallelism composed
2182 with per-module FSDP) — copy per-parameter, distributing each plain
2183 tensor onto its local shard. The default ``load_state_dict`` would
2184 recurse into the DTensor child and hit the unregistered DTensor
2185 ``copy_`` ("Operator copy_ does not contain parallel layout infer
2186 func").
2188 Args:
2189 valid_sd: Fully-qualified name → plain tensor, already shape-checked.
2190 """
2191 if isinstance(self.model, HSDPModule):
2192 self.model.load_state_dict(valid_sd, strict=False)
2193 return
2194 if not any(isinstance(p, DTensor) for _, p in self.model.named_parameters()):
2195 self.model.load_state_dict(valid_sd, strict=False)
2196 return
2197 targets: Dict[str, Any] = dict(self.model.named_parameters())
2198 targets.update(dict(self.model.named_buffers()))
2199 with platform.no_grad():
2200 for key, val in valid_sd.items():
2201 target = targets.get(key)
2202 if target is None:
2203 continue
2204 if isinstance(target, DTensor):
2205 val = _resolve_local_tensor(key, val, target)
2206 platform.load_into_param(target, val)
2208 def _load_hf_safetensors(self, weights_path: str, adapter_cls) -> None:
2209 """Load checkpoint safetensors via spec's ``state_dict_adapter``; drop shape mismatches."""
2210 # Cast loaded params down to the checkpoint's advertised dtype so the
2211 # fp32 master matches what forward consumes.
2212 load_dtype = self._resolve_hf_load_dtype(weights_path)
2213 adapter = adapter_cls()
2214 hf_sd = adapter.load_hf_state_dict(
2215 weights_path, self.model.config, dtype=load_dtype,
2216 )
2217 # Apply model-provided TP load transforms: slice the full checkpoint
2218 # weight onto this rank's shard for parameters the parallelize plan
2219 # sliced manually as plain (non-DTensor) tensors — e.g. Qwen3.5 GatedDeltaNet
2220 # ``conv1d`` / ``dt_bias`` / ``A_log`` under TP. The model is built on
2221 # meta and sliced before load, so without this the size-mismatched full
2222 # weight would be dropped (the shard then trains from random init).
2223 transform_fn = getattr(self.spec, "tp_load_transform_fn", None)
2224 if transform_fn is not None:
2225 for key, fn in transform_fn(self.model, self.mesh, self.args).items():
2226 if key in hf_sd:
2227 hf_sd[key] = fn(hf_sd[key])
2228 valid_sd, dropped, missing, unexpected = self._validate_hf_state_dict(hf_sd)
2229 if dropped:
2230 logger.warning(
2231 "Dropped %d keys due to shape mismatch (first 5: %s)",
2232 len(dropped), dropped[:5],
2233 )
2234 # Derive missing/unexpected ourselves — ``HSDPModule.load_state_dict``
2235 # returns ``None``.
2236 self._load_validated_state_dict(valid_sd)
2237 model_name = self.args.model.name
2238 logger.info_rank0(
2239 "HF (%s) load: %d tensors into hyper model",
2240 model_name, len(valid_sd),
2241 )
2242 if missing:
2243 logger.warning(
2244 "Missing (randomly initialised): %d keys, e.g. %s ...",
2245 len(missing), missing[:5],
2246 )
2247 if unexpected:
2248 logger.warning(
2249 "Unexpected (ignored): %d keys, e.g. %s ...",
2250 len(unexpected), unexpected[:5],
2251 )
2253 def _resolve_hf_load_dtype(self, weights_path: str):
2254 """Resolve the dtype to cast loaded checkpoint tensors to."""
2255 dtype_map = {
2256 'bfloat16': torch.bfloat16, 'bf16': torch.bfloat16,
2257 'float16': torch.float16, 'fp16': torch.float16,
2258 'float32': torch.float32, 'fp32': torch.float32,
2259 }
2260 cfg_dtype = (
2261 getattr(self.model.config, 'dtype', None)
2262 or getattr(self.model.config, 'torch_dtype', None)
2263 )
2264 if cfg_dtype is None:
2265 cfg_json = os.path.join(weights_path, 'config.json')
2266 if os.path.isfile(cfg_json):
2267 try:
2268 with open(cfg_json, 'r', encoding='utf-8') as f:
2269 cfg = json.load(f)
2270 cfg_dtype = cfg.get('dtype') or cfg.get('torch_dtype')
2271 except (OSError, json.JSONDecodeError):
2272 cfg_dtype = None
2273 if isinstance(cfg_dtype, str):
2274 return dtype_map.get(cfg_dtype)
2275 if isinstance(cfg_dtype, torch.dtype):
2276 return cfg_dtype
2277 return None
2279 def _validate_hf_state_dict(self, hf_sd: dict):
2280 """Strip wrapper segments and drop tensors whose shape differs from the model.
2282 Pre-validate shapes: ``load_state_dict`` aborts on the first mismatch
2283 and leaves later keys un-loaded.
2285 Returns:
2286 ``(valid_sd, dropped, missing, unexpected)``.
2287 """
2288 # Strip activation-checkpoint wrapper segments so loader keys match
2289 # ``named_parameters`` paths. The root module's parameter walk bypasses
2290 # each wrapper's own name-stripping override, so the segment leaks into
2291 # the FQN here. Covers the torch-native checkpoint_wrapper
2292 # (``_checkpoint_wrapped_module``), the hyper torch activation wrapper
2293 # (``_swap_wrapped_module``), and the hyper MindSpore activation wrapper
2294 # (``_ckpt_wrapped_module``); stripping an absent segment is a no-op.
2295 wrapper_segments = (
2296 "._checkpoint_wrapped_module",
2297 "._swap_wrapped_module",
2298 "._ckpt_wrapped_module",
2299 )
2300 def _strip(k: str) -> str:
2301 for s in wrapper_segments:
2302 k = k.replace(s, "")
2303 return k
2304 logical_to_real = {}
2305 real_to_param = {}
2306 for name, param in self.model.named_parameters():
2307 logical_to_real[_strip(name)] = name
2308 real_to_param[name] = param
2309 valid_sd: dict = {}
2310 dropped: list = []
2311 for hf_name, hf_tensor in hf_sd.items():
2312 real_name = logical_to_real.get(hf_name)
2313 if real_name is None:
2314 continue
2315 tgt = tuple(real_to_param[real_name].shape)
2316 src = tuple(hf_tensor.shape)
2317 if src == tgt:
2318 valid_sd[real_name] = hf_tensor
2319 else:
2320 dropped.append((real_name, src, tgt))
2321 param_names = set(real_to_param.keys())
2322 loaded_names = set(valid_sd.keys())
2323 missing = sorted(param_names - loaded_names)
2324 unexpected = sorted(loaded_names - param_names)
2325 return valid_sd, dropped, missing, unexpected
2327 def _load_hyper_dcp(self, weights_path: str) -> None:
2328 """Load weights from hyper's own DCP checkpoint format."""
2329 model_sd = self.model.state_dict()
2330 dcp_load(model_sd, checkpoint_id=weights_path, use_collectives=False)
2331 self.model.load_state_dict(model_sd)
2333 def _load_single_file(self, weights_path: str) -> None:
2334 """Load weights from a single ``.pt`` / ``.safetensors`` / ``.bin`` file."""
2335 sd = torch.load(weights_path, map_location="cpu", weights_only=True)
2336 missing, unexpected = self.model.load_state_dict(sd, strict=False)
2337 if missing:
2338 logger.warning("Missing keys when loading weights: %s", missing)
2339 if unexpected:
2340 logger.warning("Unexpected keys when loading weights: %s", unexpected)
2342 def _maybe_toggle_reshard(self, micro_step: int, num_micro_steps: int):
2343 """Toggle FSDP reshard_after_backward for gradient accumulation optimization.
2345 During gradient accumulation, skip resharding between micro-steps to avoid
2346 redundant all-gather. Only reshard after the last micro-step.
2347 """
2348 if not isinstance(self.model, HSDPModule) or num_micro_steps <= 1:
2349 return
2350 if micro_step == 0:
2351 self.model.set_reshard_after_backward(False)
2352 elif micro_step == num_micro_steps - 1:
2353 self.model.set_reshard_after_backward(True)