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