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