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