Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / base.py: 35%
673 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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 references:
18- BaseTrainer (``): composition pattern, 13-step init
19- hyper ``fsdp_demo.py``: FSDP wrapping via ``model.layers`` iteration
20- hyper st tests: TP → CP → AC → FSDP composition order
22Subclasses (LLMTrainer, DiTTrainer, VLMTrainer) use the composition pattern:
23instantiate ``BaseTrainer`` and call ``_build_*`` methods selectively, overriding
24or skipping steps as needed.
25"""
26import json
27import logging
28import math
29import os
30from contextlib import nullcontext
31from typing import TYPE_CHECKING, Any, Dict, Optional
33import torch
34from torch.utils.data import Dataset, DistributedSampler
36from hyper_parallel import (
37 get_platform,
38 init_empty_weights,
39 init_process_group,
40 destroy_process_group,
41 hsdp_sync_stream,
42 SkipDTensorDispatch,
43 HSDPModule,
44)
45from hyper_parallel.core.distributed_checkpoint import load as dcp_load
46from hyper_parallel.core.dtensor.dtensor import DTensor
47from hyper_parallel.core.fully_shard.hsdp_utils import GroupInfo
48from hyper_parallel.core.utils import clip_grad_norm_
49from hyper_parallel.models.spec.registry import get_spec
50from hyper_parallel.trainer.parallel_dims import ParallelDims
51from hyper_parallel.trainer.utils.loss import count_loss_token, mean_global_loss
52from hyper_parallel.trainer.callbacks.base import (
53 LoggingCallback,
54 CheckpointCallback,
55 SafetensorsExportCallback,
56 EvalCallback,
57 ProfilerCallback,
58 WandbCallback,
59 ProgressCallback,
60 MoEMonitorCallback,
61 GradientHealthCallback,
62 GCCallback,
63 TensorBoardCallback,
64 MemoryMonitorCallback,
65)
67if TYPE_CHECKING:
68 # Type-only imports — never executed at runtime, so the platform-agnostic
69 # rule ("no torch/mindspore in trainer code") is preserved. Same pattern
70 # as
71 from torch import nn
72 from torch.optim import Optimizer
73 from torch.optim.lr_scheduler import LRScheduler
74 from torch.utils.data import DataLoader
75 from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
77platform = get_platform()
78logger = logging.getLogger(__name__)
81class TrainerState:
82 """Mutable training state shared across callbacks.
84 Attributes:
85 global_step: Current training step (update count).
86 epoch: Current epoch index.
87 max_steps: Total number of training steps.
88 """
90 def __init__(self, max_steps: int = 0):
91 self.global_step: int = 0
92 self.epoch: int = 0
93 self.max_steps: int = max_steps
94 self.log_history: list = []
97class BaseTrainer:
98 """Composable training skeleton.
100 Provides 13 ``_build_*`` methods that subclasses can call, override, or skip.
101 The default ``_build_parallelized_model`` applies TP → CP → AC → FSDP by
102 iterating ``model.layers`` — matching hyper's own ``fsdp_demo.py`` style.
104 Args:
105 args: Training configuration (typically parsed from YAML).
106 """
108 # PEP 526 annotations — populated by ``_build_*``; ``None`` until built.
109 model: Optional["nn.Module"] = None
110 optimizer: Optional["Optimizer"] = None
111 lr_scheduler: Optional["LRScheduler"] = None
112 train_dataloader: Optional["DataLoader"] = None
113 mesh: Optional["DeviceMesh"] = None
114 device = None
115 parallel_dims = None
116 _dp_group_info = None
117 tokenizer = None
118 processor = None
119 data_transform = None
120 train_dataset = None
121 collate_fn = None
122 _grad_accum = None
123 model_fwd_context = None
124 grad_scaler = None
125 model_bwd_context = None
126 logging_callback = None
127 checkpoint_callback = None
128 hf_export_callback = None
129 eval_callback = None
130 profiler_callback = None
131 wandb_callback = None
132 tensorboard_callback = None
133 progress_callback = None
134 moe_monitor_callback = None
135 gradient_health_callback = None
136 memory_monitor_callback = None
137 gc_callback = None
138 user_callbacks: list = None
139 sampler = None
140 _last_global_tokens = None
142 def __init__(self, args):
143 # Only early-bound fields live here; the rest is built via
144 # ``_build_*`` methods invoked by the subclass.
145 self.args = args
146 self.spec = get_spec(args.model.name)
147 self.state = TrainerState(max_steps=getattr(args.train, 'max_steps', 0))
149 # ------------------------------------------------------------------
150 # 13 overridable _build_* methods
151 # ------------------------------------------------------------------
153 def _setup(self):
154 """Step 1: Initialize distributed environment, device mesh, and seed.
156 Calls hyper's own ``init_process_group`` and ``init_device_mesh``.
157 Mesh shape is derived from ``args.parallel`` (dp, tp, cp, pp, ep).
158 """
159 backend = getattr(self.args.train, 'comm_backend', None)
160 init_process_group(backend=backend)
162 local_rank = getattr(self.args.train, 'local_rank', 0)
163 device_type = platform.device_type() # "npu" or "cuda"
164 # Use platform.device(idx) — backend-agnostic.
165 self.device = platform.device(local_rank)
166 device_handle = platform.get_device_handle(device_type)
167 device_handle.set_device(local_rank)
169 # Build & validate parallel dims in one place (fail-fast).
171 self.parallel_dims = ParallelDims.from_config(
172 self.args.train.accelerator, world_size=platform.get_world_size(),
173 )
174 logger.info_rank0("ParallelDims: %s", self.parallel_dims.summary())
175 self.mesh = self.parallel_dims.build_mesh(platform.device_type())
177 # Build DP group_info for trainer-level all_reduce (loss/token sync).
178 # Uses hyper's GroupInfo + mesh.get_group (platform-agnostic).
180 dp_group = self._get_combined_dp_group()
181 dp_size = self.parallel_dims.dp_size
182 self._dp_group_info = GroupInfo(
183 group_name="trainer_dp", group=dp_group, rank_size=dp_size,
184 )
186 debug_cfg = getattr(self.args.train, 'debug', None)
187 seed = getattr(self.args.train, 'seed', 42)
188 platform.manual_seed(seed)
189 # Seed device-side RNG explicitly: ``platform.manual_seed`` only
190 # covers CPU.
191 try:
192 handle = platform.get_device_handle(device_type)
193 if hasattr(handle, "manual_seed_all"):
194 handle.manual_seed_all(seed)
195 elif hasattr(handle, "manual_seed"):
196 handle.manual_seed(seed)
197 except Exception as exc: # pylint: disable=W0718
198 logger.warning("Device-side seed init skipped: %s", exc)
200 if debug_cfg is not None and getattr(debug_cfg, 'deterministic', False):
201 warn_only = getattr(debug_cfg, 'deterministic_warn_only', False)
202 torch.use_deterministic_algorithms(True, warn_only=warn_only)
203 torch.backends.cudnn.deterministic = True
204 torch.backends.cudnn.benchmark = False
205 logger.info_rank0("Deterministic algorithms enabled (warn_only=%s)", warn_only)
207 logger.info_rank0(
208 "Setup complete: rank=%d, world_size=%d, mesh=%s",
209 platform.get_rank(), platform.get_world_size(),
210 self.mesh.mesh_dim_names,
211 )
212 logger.info_rank0(
213 "Config: data.type=%s, model.name=%s, model.num_hidden_layers=%s, "
214 "init_device=%s, max_steps=%d, global_bs=%d",
215 getattr(self.args.data, 'type', '?'),
216 getattr(self.args.model, 'name', '?'),
217 getattr(self.args.model, 'num_hidden_layers', '?'),
218 getattr(self.args.train, 'init_device', '?'),
219 self.state.max_steps,
220 getattr(self.args.train, 'global_batch_size', '?'),
221 )
223 def _build_model(self):
224 """Step 2: Construct model via ``spec.build_model_fn``.
226 The model is a plain ``nn.Module`` at this point — not yet parallelized.
227 When ``args.runtime.init_device == "meta"``, the model is constructed on
228 the meta device (no memory allocated) and real weights are loaded after
229 FSDP sharding via ``_load_weights_after_parallel``.
230 """
231 init_device = getattr(self.args.train, 'init_device', 'meta')
232 # Meta-device init: each rank materialises only its own shard
233 # post-FSDP — pre-trained weights via DCP, otherwise random init.
234 if init_device == "meta":
236 with init_empty_weights():
237 self.model = self.spec.build_model_fn(self.args)
238 logger.info_rank0(
239 "Model built on meta device (no memory allocated): %s",
240 type(self.model).__name__,
241 )
242 else:
243 self.model = self.spec.build_model_fn(self.args)
244 logger.info_rank0("Model built on %s: %s", init_device, type(self.model).__name__)
246 # Cross-check parallel degrees against the actual model hyperparams
247 # (heads%tp, kv_heads%tp, num_experts%ep, seq_len%(cp*tp)).
248 # Fails fast here instead of crashing inside parallelize_module.
249 seq_len = getattr(self.args.data, 'max_seq_len', None)
250 self.parallel_dims.validate_against_model(self.model, seq_len=seq_len)
252 def _freeze_model(self):
253 """Step 3: Freeze specified modules (optional)."""
254 freeze_modules = getattr(self.args.model, 'freeze_modules', None)
255 if not freeze_modules:
256 return
257 for name, param in self.model.named_parameters():
258 if any(pattern in name for pattern in freeze_modules):
259 param.requires_grad_(False)
261 def _build_model_assets(self):
262 """Step 4: Build tokenizer, processor, chat_template.
264 Default: no-op. LLMTrainer overrides to build tokenizer + chat_template.
265 VLMTrainer overrides to build processor.
266 """
267 self.tokenizer = None
268 self.processor = None
270 def _build_data_transform(self):
271 """Step 5: Build data preprocessing transform.
273 Default: identity transform. LLMTrainer overrides for tokenization.
274 """
275 self.data_transform = None
277 def _build_dataset(self):
278 """Step 6: Build training dataset.
280 Supports:
281 - ``data.type = "dummy"``: random tokens for validation
282 - ``data.type = "hf_datasets"``: HuggingFace datasets
283 - ``data.type = "megatron_indexed"``: Megatron .bin/.idx format
285 Subclass can override for custom dataset logic.
286 """
288 data_type = getattr(self.args.data, 'type', 'dummy')
289 seq_len = getattr(self.args.data, 'max_seq_len', 2048)
291 if data_type == "dummy":
292 vocab_size = getattr(self.model, 'config', None)
293 vocab_size = vocab_size.vocab_size if vocab_size else 32000
294 total_samples = self.state.max_steps * getattr(self.args.train, 'global_batch_size', 8)
296 class DummyDataset(Dataset):
297 """Deterministic random token dataset for FSDP validation.
299 Each sample's content is fixed by its index (seeded by
300 ``base_seed + idx``), so the same index always produces the
301 same tokens regardless of access order or DP configuration.
302 """
303 def __init__(self, num_samples, seq_length, vocab, base_seed=42):
304 self.num_samples = num_samples
305 self.seq_length = seq_length
306 self.vocab = vocab
307 self.base_seed = base_seed
309 def __len__(self):
310 return self.num_samples
312 def __getitem__(self, idx):
313 g = torch.Generator().manual_seed(self.base_seed + idx)
314 input_ids = torch.randint(
315 0, self.vocab, (self.seq_length,), generator=g,
316 )
317 return {"input_ids": input_ids, "labels": input_ids.clone()}
319 self.train_dataset = DummyDataset(total_samples, seq_len, vocab_size)
320 logger.info_rank0("Dummy dataset created: %d samples, seq_len=%d", total_samples, seq_len)
322 elif data_type == "hf_datasets":
323 from datasets import load_dataset # pylint: disable=C0415 # optional dep
324 train_path = self.args.data.train_path
325 streaming = getattr(self.args.data, 'streaming', False)
326 if streaming:
327 # ``DistributedSampler`` requires ``__len__``, which
328 # ``IterableDataset`` lacks; reject loudly until a
329 # sampler-less streaming path is wired.
330 raise NotImplementedError(
331 "data.streaming=True not yet wired. The current "
332 "_build_dataloader uses DistributedSampler which requires "
333 "len(dataset); IterableDataset has no __len__. "
334 "Use data.streaming=False (map-style) for now, or "
335 "subclass _build_dataset + _build_dataloader to emit an "
336 "IterableDataset that self-shards via dp_rank/dp_size."
337 )
338 ds = load_dataset(train_path, split="train", streaming=False)
339 if self.data_transform:
340 ds = ds.map(
341 self.data_transform, remove_columns=ds.column_names,
342 )
343 self.train_dataset = ds
344 logger.info_rank0("HF dataset loaded: %s (map-style)", train_path)
346 else:
347 raise ValueError(f"Unknown data type: {data_type}. Supported: dummy, hf_datasets")
349 def _build_collate_fn(self):
350 """Step 7: Build data collator.
352 Default: pads input_ids and labels to max length in the batch.
353 """
355 def _default_collate(batch):
356 """Simple padding collator."""
357 max_len = max(item["input_ids"].size(0) for item in batch)
358 input_ids_list = []
359 labels_list = []
360 for item in batch:
361 pad_len = max_len - item["input_ids"].size(0)
362 input_ids_list.append(
363 torch.nn.functional.pad(item["input_ids"], (0, pad_len), value=0)
364 )
365 labels_list.append(
366 torch.nn.functional.pad(item["labels"], (0, pad_len), value=-100)
367 )
368 return {
369 "input_ids": torch.stack(input_ids_list),
370 "labels": torch.stack(labels_list),
371 }
373 self.collate_fn = _default_collate
375 def _build_dataloader(self):
376 """Step 8: Build distributed stateful dataloader.
378 Uses ``torchdata.stateful_dataloader.StatefulDataLoader`` so that
379 iterator position is checkpointable — enabling exact resume after
380 restart (matching ).
382 Each ``next()`` call yields a list of micro-batches (for gradient
383 accumulation).
384 """
385 from torchdata.stateful_dataloader import StatefulDataLoader # pylint: disable=C0415 # optional dep
387 micro_bs = getattr(self.args.train, 'micro_batch_size', 1)
389 # Sampler uses DP rank/size — TP/CP/PP/EP peers share data.
390 dp_size = self.parallel_dims.dp_size
391 non_dp = self.parallel_dims.non_dp_size
392 global_rank = platform.get_rank()
393 dp_rank = global_rank // non_dp if non_dp > 1 else global_rank
395 shuffle = getattr(self.args.data, "shuffle", True)
396 sampler_seed = getattr(self.args.train, 'seed', 0)
398 self.sampler = DistributedSampler(
399 self.train_dataset,
400 num_replicas=dp_size,
401 rank=dp_rank,
402 shuffle=shuffle,
403 seed=sampler_seed,
404 drop_last=True,
405 )
407 # StatefulDataLoader supports state_dict() / load_state_dict()
408 # for checkpoint resume (torchdata API, used by + ).
409 num_workers = getattr(self.args.data, 'num_workers', 0)
410 prefetch_factor = getattr(self.args.data, 'prefetch_factor', None)
411 pin_memory = getattr(self.args.data, 'pin_memory', True)
412 loader_kwargs = {
413 "batch_size": micro_bs,
414 "sampler": self.sampler,
415 "collate_fn": self.collate_fn,
416 "num_workers": num_workers,
417 "pin_memory": pin_memory,
418 "drop_last": True,
419 }
420 # prefetch_factor is only accepted when num_workers > 0
421 if num_workers > 0 and prefetch_factor is not None:
422 loader_kwargs["prefetch_factor"] = prefetch_factor
423 self.train_dataloader = StatefulDataLoader(
424 self.train_dataset, **loader_kwargs,
425 )
427 # Use dp_size (not world_size) — TP/CP/PP ranks share data, not split it.
428 self._grad_accum = max(
429 getattr(self.args.train, 'global_batch_size', micro_bs) // (micro_bs * dp_size),
430 1,
431 )
433 logger.info_rank0(
434 "Dataloader built: micro_bs=%d, grad_accum=%d, dataset_size=%d",
435 micro_bs, self._grad_accum, len(self.train_dataset),
436 )
438 def _build_parallelized_model(self):
439 """Step 9: Apply parallel strategies to the model.
441 Each model owns its full parallelize pipeline in
442 ``models/<name>/parallelize.py`` (convention) and
443 registers it via ``ModelSpec.parallelize_fn``. There is no shared
444 "default" template — model-specific TP/EP/CP/AC/FSDP/Prefetch
445 composition lives next to the model that needs it.
446 """
447 if self.spec.parallelize_fn is None:
448 raise ValueError(
449 f"Model '{self.spec.name}' has no ``parallelize_fn`` registered "
450 f"on its ModelSpec. Each model must own its parallelize "
451 f"pipeline in models/<name>/parallelize.py."
452 )
453 self.model = self.spec.parallelize_fn(self.model, self.mesh, self.args)
454 self._post_parallelize()
456 def _post_parallelize(self):
457 """Common steps after parallelization (materialize weights + train mode).
459 Order when ``init_device == "meta"`` and ``weights_path`` is set:
461 1. Run ``_materialize_and_init_shards`` first — this calls
462 ``model.to_empty(device=...)`` + kaiming / zero init for every
463 parameter. That is the **baseline** state so no param stays on
464 meta (which would trip ``HSDPState._validate_no_meta_params``).
465 2. Then ``_load_weights`` copies the upstream checkpoint on top.
466 Every key that matches overwrites the random init; anything
467 missing in the checkpoint stays with its kaiming / zero init.
469 This pattern handles partial checkpoints cleanly — e.g. hyper's
470 current Qwen3-VL-MoE model lacks ``q_norm`` / ``k_norm`` /
471 ``pos_embed`` / ``deepstack_merger_list``; those stay random
472 while all other weights come from the pretrained checkpoint.
473 """
474 init_device = getattr(self.args.train, 'init_device', 'meta')
475 weights_path = getattr(self.args.model, 'weights_path', None)
476 if init_device == "meta":
477 # Always materialize first (random init baseline) so no param
478 # stays on meta — then overlay the checkpoint.
479 self._materialize_and_init_shards()
480 if weights_path:
481 self._load_weights(weights_path)
482 elif weights_path:
483 self._load_weights(weights_path)
484 # Mixed-precision storage policy: trainable params keep an fp32
485 # master; frozen params keep their loaded low-precision dtype so the
486 # forward stays uniform-bf16 within their FSDP unit.
487 self._maybe_downcast_frozen_params()
488 self._maybe_upcast_trainable_params()
489 self.model.train()
491 def _maybe_downcast_frozen_params(self) -> None:
492 """Maybe downcast frozen params (internal)."""
493 freeze_modules = getattr(self.args.model, 'freeze_modules', None)
494 if not freeze_modules:
495 return
496 mp_cfg = getattr(self.args.train, 'mixed_precision', None)
497 if mp_cfg is None or not getattr(mp_cfg, 'enabled', False):
498 return
500 target_dtype = {
501 'bfloat16': torch.bfloat16,
502 'bf16': torch.bfloat16,
503 'float16': torch.float16,
504 'fp16': torch.float16,
505 }.get(getattr(mp_cfg, 'param_dtype', 'bfloat16'))
506 if target_dtype is None:
507 return
508 n_cast = 0
509 for name, param in self.model.named_parameters():
510 if not any(pat in name for pat in freeze_modules):
511 continue
512 if param.requires_grad:
513 continue
514 local = param.data
515 if hasattr(local, 'to_local'):
516 local = local.to_local()
517 if local.dtype == target_dtype:
518 continue
519 new_local = local.to(target_dtype)
520 # DTensor: rebuild the global view via from_local with same placements.
521 if hasattr(param.data, 'to_local'):
523 if isinstance(param.data, DTensor):
524 param.data = DTensor.from_local(
525 new_local,
526 device_mesh=param.data.device_mesh,
527 placements=param.data.placements,
528 )
529 else:
530 param.data = new_local
531 else:
532 param.data = new_local
533 n_cast += 1
534 logger.info_rank0(
535 "Post-load: cast %d frozen params to %s",
536 n_cast, target_dtype,
537 )
539 def _maybe_upcast_trainable_params(self) -> None:
540 """Upcast trainable params to ``float32``.
542 Implements the standard mixed-precision pattern (fp32 master weight
543 + low-precision forward): trainable params loaded from a bf16
544 checkpoint would otherwise stay in bf16, and Adam moment estimates
545 accumulating at bf16's 7-bit mantissa diverge noticeably after only
546 a handful of optimizer steps.
548 Runs AFTER ``_maybe_downcast_frozen_params`` so frozen params keep
549 their bf16 storage and only trainable params get the fp32 master.
550 """
551 mp_cfg = getattr(self.args.train, 'mixed_precision', None)
552 if mp_cfg is None or not getattr(mp_cfg, 'enabled', False):
553 return
555 n_cast = 0
556 for _, param in self.model.named_parameters():
557 if not param.requires_grad:
558 continue
559 local = param.data
560 if hasattr(local, 'to_local'):
561 local = local.to_local()
562 if local.dtype == torch.float32:
563 continue
564 new_local = local.to(torch.float32)
565 if hasattr(param.data, 'to_local'):
567 if isinstance(param.data, DTensor):
568 param.data = DTensor.from_local(
569 new_local,
570 device_mesh=param.data.device_mesh,
571 placements=param.data.placements,
572 )
573 else:
574 param.data = new_local
575 else:
576 param.data = new_local
577 n_cast += 1
578 logger.info_rank0(
579 "Post-load: upcast %d trainable params to float32", n_cast,
580 )
582 def _build_optimizer(self):
583 """Step 10: Build optimizer. Must be called AFTER ``_build_parallelized_model``.
585 After FSDP, parameters are DTensor shards — optimizer operates on local shards.
586 Optimizer must be created after ``fully_shard``.
587 """
588 lr = getattr(self.args.train.optimizer, 'lr', 1e-4)
589 weight_decay = getattr(self.args.train.optimizer, 'weight_decay', 0.01)
591 # bias / LayerNorm / RMSNorm go to no-decay; grouping matters even
592 # at wd=0 — foreach Adam reduction order differs per group on NPU.
593 decay_keywords = ("bias", "layernorm", "norm", "rmsnorm")
595 def _is_no_decay(name: str) -> bool:
596 lname = name.lower()
597 return any(kw in lname for kw in decay_keywords)
599 decay_params = []
600 no_decay_params = []
601 seen_ids = set()
602 for n, p in self.model.named_parameters():
603 if not p.requires_grad:
604 continue
605 # Dedup tied params (same nn.Parameter shared across modules).
606 if id(p) in seen_ids:
607 continue
608 seen_ids.add(id(p))
609 if _is_no_decay(n):
610 no_decay_params.append(p)
611 else:
612 decay_params.append(p)
614 param_groups = [
615 {"params": decay_params, "weight_decay": weight_decay},
616 {"params": no_decay_params, "weight_decay": 0.0},
617 ]
618 adam_eps = getattr(self.args.train.optimizer, 'eps', 1e-8)
619 adam_betas = getattr(self.args.train.optimizer, 'betas', (0.9, 0.999))
620 adam_foreach = getattr(self.args.train.optimizer, 'foreach', None)
621 self.optimizer = torch.optim.AdamW(
622 param_groups,
623 lr=lr,
624 betas=adam_betas,
625 eps=adam_eps,
626 foreach=adam_foreach,
627 )
628 logger.info_rank0(
629 "Optimizer: AdamW lr=%.2e wd=%.3g decay_params=%d no_decay_params=%d",
630 lr, weight_decay, len(decay_params), len(no_decay_params),
631 )
633 def _build_lr_scheduler(self):
634 """Step 11: Build learning rate scheduler.
636 Supports cosine decay with warmup. Falls back to constant LR if
637 warmup_ratio is 0 and decay_style is 'constant'.
638 """
640 total_steps = self.state.max_steps
641 warmup_ratio = getattr(self.args.train.optimizer, 'lr_warmup_ratio', 0.0)
642 # ``ceil`` matches the standard warmup convention so a fractional
643 # ``warmup_ratio * max_steps`` rounds up to the next full step.
644 warmup_steps = math.ceil(total_steps * warmup_ratio)
645 decay_style = getattr(self.args.train.optimizer, 'lr_decay_style', 'cosine')
646 lr_min = getattr(self.args.train.optimizer, 'lr_min', 0.0)
647 lr_max = getattr(self.args.train.optimizer, 'lr', 1e-4)
649 def _lr_lambda(current_step):
650 if current_step < warmup_steps:
651 return float(current_step) / float(max(1, warmup_steps))
652 if decay_style == 'constant':
653 return 1.0
654 # Cosine decay
655 progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
656 cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
657 min_ratio = lr_min / lr_max if lr_max > 0 else 0.0
658 return min_ratio + (1.0 - min_ratio) * cosine_decay
660 self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, _lr_lambda)
661 logger.info_rank0(
662 "LR scheduler: %s, warmup_steps=%d/%d, lr=%.2e→%.2e",
663 decay_style, warmup_steps, total_steps, lr_max, lr_min,
664 )
666 def _build_training_context(self):
667 """Step 12: Build forward/backward context managers.
669 Delegates AMP context to ``_make_amp_context`` — the single place that
670 knows about backend-specific autocast. MindSpore subclass should
671 override ``_make_amp_context`` only.
672 """
675 mp_cfg = getattr(self.args.train, 'mixed_precision', None)
676 # FSDP2 mp_policy already runs forward/backward in low precision;
677 # stacking ``torch.autocast`` on top would re-promote LayerNorm /
678 # softmax / log_softmax to fp32 and break the uniform-bf16 forward.
679 self.model_fwd_context = nullcontext()
680 self.grad_scaler = None
681 if mp_cfg and getattr(mp_cfg, 'enabled', False):
682 param_dtype_str = getattr(mp_cfg, 'param_dtype', 'bfloat16')
683 logger.info_rank0(
684 "Mixed precision via FSDP2 mp_policy: dtype=%s on %s "
685 "(autocast disabled — pure low-precision forward under mp_policy)",
686 param_dtype_str, platform.device_type(),
687 )
689 self.model_bwd_context = nullcontext()
691 @staticmethod
692 def _make_amp_context(param_dtype_str: str):
693 """Build the AMP forward context. Backend-specific.
695 This is the SOLE method in the trainer that touches the backend AMP
696 API directly. Override this single method to add MindSpore support
697 (``ms.amp.auto_mixed_precision``). Default uses ``torch.autocast``.
698 """
699 dtype = getattr(torch, param_dtype_str, torch.bfloat16)
700 return torch.autocast(platform.device_type(), dtype=dtype)
702 def _init_callbacks(self):
703 """Step 13: Initialize callbacks (explicit mode).
705 Each callback is a named field — engineer sees all callbacks and their
706 order in ``on_step_end`` at a glance. Add/remove/reorder = change one line.
707 """
708 self.logging_callback = LoggingCallback(self)
709 self.checkpoint_callback = CheckpointCallback(self)
710 self.hf_export_callback = SafetensorsExportCallback(self)
711 self.eval_callback = EvalCallback(self)
712 self.profiler_callback = ProfilerCallback(self)
713 self.wandb_callback = WandbCallback(self)
714 self.tensorboard_callback = TensorBoardCallback(self)
715 self.progress_callback = ProgressCallback(self)
716 self.moe_monitor_callback = MoEMonitorCallback(self)
717 # Health + operability (no-ops unless enabled in cfg.train.debug / .memory_monitor).
718 self.gradient_health_callback = GradientHealthCallback(self)
719 self.memory_monitor_callback = MemoryMonitorCallback(self)
720 self.gc_callback = GCCallback(self)
721 # ``user_callbacks`` lets external code append extra Callback instances
722 # (e.g. domain-specific monitors) without editing this method. They get
723 # the same lifecycle dispatch as built-ins.
724 self.user_callbacks: list = []
725 logger.info_rank0(
726 "Callbacks initialized: logging, checkpoint, hf_export, eval, "
727 "profiler, wandb, tensorboard, progress, moe_monitor, "
728 "gradient_health, memory_monitor, gc"
729 )
731 # ------------------------------------------------------------------
732 # Public API: external callback registration
733 # ------------------------------------------------------------------
735 def add_callback(self, callback) -> None:
736 """Register an extra ``Callback`` to receive every lifecycle event.
738 Use this to plug domain-specific monitors (custom metric sinks,
739 in-house experiment trackers, RL reward loggers) without editing
740 the trainer. Built-in callbacks always run first; user callbacks
741 run in registration order so a later user callback can read state
742 the earlier ones updated.
743 """
744 self.user_callbacks.append(callback)
745 logger.info_rank0(
746 "User callback registered: %s", type(callback).__name__,
747 )
749 # ------------------------------------------------------------------
750 # Callback dispatch (explicit mode)
751 # ------------------------------------------------------------------
753 def _builtin_callbacks(self) -> list:
754 """Return built-in callbacks in fixed dispatch order.
756 Centralised so every dispatcher iterates the same list — adding a
757 callback only needs an entry here plus a named field in
758 ``_init_callbacks`` (no per-event copy/paste).
759 """
760 return [
761 self.logging_callback,
762 self.eval_callback,
763 self.profiler_callback,
764 self.wandb_callback,
765 self.tensorboard_callback,
766 self.progress_callback,
767 self.checkpoint_callback,
768 self.hf_export_callback,
769 self.moe_monitor_callback,
770 self.gradient_health_callback,
771 self.memory_monitor_callback,
772 self.gc_callback,
773 ]
775 def _all_callbacks(self) -> list:
776 """Built-in callbacks followed by user-registered ones."""
777 return self._builtin_callbacks() + list(self.user_callbacks)
779 def on_init_end(self):
780 """Dispatch one-shot ``on_init_end`` after every ``_build_*`` ran.
782 Fired by the subclass at the end of its own ``__init__`` (see
783 ``LLMTrainer.__init__``); ``BaseTrainer.train()`` does NOT call it
784 because BaseTrainer instances are sometimes wrapped (composition
785 pattern) and the wrapper owns the init lifecycle.
786 """
787 for cb in self._all_callbacks():
788 cb.on_init_end(self.state)
790 def on_train_begin(self):
791 """Dispatch on_train_begin to all callbacks."""
792 # Memory monitor first so it captures the truly-initial peak.
793 self.memory_monitor_callback.on_train_begin(self.state)
794 self.moe_monitor_callback.on_train_begin(self.state)
795 self.profiler_callback.on_train_begin(self.state)
796 self.wandb_callback.on_train_begin(self.state)
797 self.tensorboard_callback.on_train_begin(self.state)
798 self.progress_callback.on_train_begin(self.state)
799 # Checkpoint runs LAST so resume sees an already-armed TB writer
800 # (it'll record the load event via dispatch_load_event).
801 self.checkpoint_callback.on_train_begin(self.state)
802 for cb in self.user_callbacks:
803 cb.on_train_begin(self.state)
805 def on_train_end(self):
806 """Dispatch on_train_end to all callbacks."""
807 self.checkpoint_callback.on_train_end(self.state)
808 self.hf_export_callback.on_train_end(self.state)
809 self.progress_callback.on_train_end(self.state)
810 self.tensorboard_callback.on_train_end(self.state)
811 self.wandb_callback.on_train_end(self.state)
812 self.profiler_callback.on_train_end(self.state)
813 for cb in self.user_callbacks:
814 cb.on_train_end(self.state)
816 def on_step_begin(self):
817 """Dispatch on_step_begin to all callbacks."""
818 self.logging_callback.on_step_begin(self.state)
819 for cb in self.user_callbacks:
820 cb.on_step_begin(self.state)
822 def on_step_end(self, loss=None, grad_norm=None):
823 """Dispatch on_step_end to all callbacks (built-ins + user)."""
824 for cb in self._all_callbacks():
825 cb.on_step_end(self.state, loss=loss, grad_norm=grad_norm)
827 def on_substep_end(self):
828 """Dispatch on_substep_end (after each micro-batch forward/backward)."""
829 self.moe_monitor_callback.on_substep_end(self.state)
830 for cb in self.user_callbacks:
831 cb.on_substep_end(self.state)
833 def on_pre_optimizer_step(self, grad_norm=None):
834 """Dispatch on_pre_optimizer_step (after grad clip, before optimizer.step)."""
835 # Health check runs FIRST so a NaN aborts before the logger misleads.
836 self.gradient_health_callback.on_pre_optimizer_step(
837 self.state, grad_norm=grad_norm,
838 )
839 self.logging_callback.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
840 self.wandb_callback.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
841 self.tensorboard_callback.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
842 for cb in self.user_callbacks:
843 cb.on_pre_optimizer_step(self.state, grad_norm=grad_norm)
845 def on_epoch_begin(self):
846 """Dispatch on_epoch_begin."""
847 for cb in self._all_callbacks():
848 cb.on_epoch_begin(self.state)
850 def on_epoch_end(self):
851 """Dispatch on_epoch_end."""
852 for cb in self._all_callbacks():
853 cb.on_epoch_end(self.state)
855 # ------------------------------------------------------------------
856 # Event fan-out (LoggingCallback / CheckpointCallback emit these)
857 # ------------------------------------------------------------------
859 def dispatch_log_event(self, metrics: dict) -> None:
860 """Forward a metrics record to every callback's ``on_log``.
862 ``LoggingCallback`` calls this so TensorBoard / W&B / external sinks
863 log the SAME numbers — single source of truth, no duplicate work.
864 """
865 for cb in self._all_callbacks():
866 cb.on_log(self.state, metrics=metrics)
868 def dispatch_save_event(self, checkpoint_dir: str) -> None:
869 """Forward a ckpt-save event to every callback's ``on_save``."""
870 for cb in self._all_callbacks():
871 cb.on_save(self.state, checkpoint_dir=checkpoint_dir)
873 def dispatch_load_event(self, checkpoint_dir: str) -> None:
874 """Forward a ckpt-load event to every callback's ``on_load``."""
875 for cb in self._all_callbacks():
876 cb.on_load(self.state, checkpoint_dir=checkpoint_dir)
878 def dispatch_evaluate_event(self, metrics: dict = None) -> None:
879 """Forward an eval-pass-complete event to every callback's ``on_evaluate``."""
880 for cb in self._all_callbacks():
881 cb.on_evaluate(self.state, metrics=metrics)
883 # ------------------------------------------------------------------
884 # Training core
885 # ------------------------------------------------------------------
887 def forward_backward_step(
888 self,
889 micro_batch: Dict[str, Any],
890 micro_batch_tokens: int,
891 global_tokens: int,
892 num_micro: int = 1,
893 ):
894 """Run forward + backward for one micro-batch.
896 Uses global token normalisation: each micro-batch's
897 loss is scaled by ``micro_tokens / global_tokens`` so that every token
898 across all ranks and all micro-batches contributes equally to the
899 gradient, regardless of DP size or grad_accum.
901 Args:
902 micro_batch: Dict of input tensors.
903 micro_batch_tokens: Non-padding token count for this micro-batch.
904 global_tokens: Total non-padding tokens across **all** ranks and
905 **all** micro-batches (computed via all-reduce).
907 Returns:
908 Tuple of (raw_loss_scalar, micro_batch_tokens) for logging.
909 """
910 # Move tensors to device. Multimodal batches may contain nested
911 # lists/tuples/dicts, so keep this recursive instead of assuming a
912 # flat LM-only batch.
913 def _to_device(value):
914 if hasattr(value, "to"):
915 return value.to(self.device, non_blocking=True)
916 if isinstance(value, dict):
917 return {k: _to_device(v) for k, v in value.items()}
918 if isinstance(value, list):
919 return [_to_device(v) for v in value]
920 if isinstance(value, tuple):
921 return tuple(_to_device(v) for v in value)
922 return value
924 micro_batch = {k: _to_device(v) for k, v in micro_batch.items()}
926 # Forward (with training context for activation offload)
927 with self.model_fwd_context:
928 outputs = self.model(**micro_batch, use_cache=False)
929 loss = outputs["loss"] if isinstance(outputs, dict) else outputs.loss
931 # TP scenario: loss may be Partial DTensor — reduce before backward
932 if hasattr(loss, 'is_partial') and loss.is_partial():
933 loss = loss.reduce_partial()
935 # Keep raw loss value for logging before scaling
936 raw_loss = loss.detach()
938 # token_weighted: ``CE_mean * (micro_tokens / global_tokens) * dp_size``
939 # — DP-size and grad_accum invariant after FSDP's ``AVG`` reduction.
940 dp_size = self.parallel_dims.dp_size
941 agg = getattr(self.args.train.optimizer, 'loss_aggregation', 'token_weighted')
942 if agg == 'rank_average':
943 # Per-rank CE mean; FSDP2 ``ReduceOp.AVG`` handles cross-rank
944 # averaging. Dividing by ``num_micro`` averages grad-accum so
945 # 1-card grad_accum=N tracks the FSDP path's effective grad.
946 scaled_loss = loss / num_micro if num_micro > 1 else loss
947 else:
948 scaled_loss = mean_global_loss(
949 loss, micro_batch_tokens, global_tokens, dp_size,
950 )
952 # Backward (with training context)
953 with self.model_bwd_context:
954 scaled_loss.backward()
956 return raw_loss, micro_batch_tokens
958 def train_step(self, data_iterator):
959 """Execute one training step with gradient accumulation.
961 Precision-aligned across different DP configurations by:
962 1. All-reducing global token count before loss scaling ()
963 2. Syncing gradients only on the last micro-batch ()
964 3. All-reducing loss weighted by token count for reporting
966 Args:
967 data_iterator: Iterator yielding lists of micro-batch dicts.
968 """
969 # Pull data first; ``StopIteration`` propagates without bumping the
970 # step counter so checkpoint dirs / log indices match the steps that
971 # actually trained.
972 micro_batches = next(data_iterator)
973 self.state.global_step += 1
974 num_micro = len(micro_batches)
976 # ---- Phase 1: count global tokens ( style) ----
977 # All-reduce BEFORE forward so every rank uses the same denominator.
979 token_counts = [count_loss_token(mb) for mb in micro_batches]
980 local_tokens = sum(token_counts)
981 if local_tokens == 0:
982 local_tokens = 1
983 global_tokens = local_tokens
984 if platform.get_world_size() > 1:
985 gt = platform.full((1,), local_tokens).to(self.device)
986 platform.all_reduce(gt, self._dp_group_info)
987 global_tokens = max(int(gt.item()), 1)
988 # Expose for callbacks (e.g. LoggingCallback throughput).
989 self._last_global_tokens = global_tokens
991 # ---- Phase 2: forward / backward with per-micro-batch sync control ----
992 total_loss_sum = 0.0 # weighted: loss * micro_tokens
993 total_loss_arith_sum = 0.0 # arithmetic sum across micro-batches
994 total_tokens_local = 0
995 for i, mb in enumerate(micro_batches):
996 is_last = i == num_micro - 1
998 # Gradient sync: only on the last micro-batch ( pattern).
999 # Before last: accumulate gradients locally, no communication.
1000 if isinstance(self.model, HSDPModule):
1001 self.model.set_requires_gradient_sync(is_last)
1002 self.model.set_is_last_backward(is_last)
1004 # FSDP reshard optimization for gradient accumulation
1005 self._maybe_toggle_reshard(i, num_micro)
1007 raw_loss, mb_tokens = self.forward_backward_step(
1008 mb, token_counts[i], global_tokens, num_micro=num_micro,
1009 )
1010 total_loss_sum += raw_loss.item() * mb_tokens
1011 total_loss_arith_sum += raw_loss.item()
1012 total_tokens_local += mb_tokens
1013 self.on_substep_end()
1015 # Wait for async gradient reduce
1016 #
1017 hsdp_sync_stream()
1019 # Gradient clipping — DTensor-aware, returns plain Tensor
1020 max_grad_norm = getattr(self.args.train.optimizer, 'max_grad_norm', 1.0)
1021 clip_fn = self.spec.clip_grad_fn or clip_grad_norm_
1022 grad_norm = clip_fn(self.model.parameters(), max_grad_norm)
1023 self.on_pre_optimizer_step(grad_norm=grad_norm.item())
1025 # Optimizer step — must be inside SkipDTensorDispatch
1026 with SkipDTensorDispatch():
1027 self.optimizer.step()
1029 if self.lr_scheduler is not None:
1030 self.lr_scheduler.step()
1032 self.optimizer.zero_grad()
1034 # ---- Phase 3: all-reduce loss for reporting ----
1035 agg = getattr(self.args.train.optimizer, 'loss_aggregation', 'token_weighted')
1036 if agg == 'token_weighted':
1037 # Global token-weighted: sum local ce_sum across DP, divide by
1038 # global token count. Equivalent to mean-loss-per-token under
1039 # post-DP gradient mean.
1040 if platform.get_world_size() > 1:
1041 ls = platform.full((1,), total_loss_sum).to(self.device)
1042 platform.all_reduce(ls, self._dp_group_info)
1043 avg_loss = ls.item() / max(global_tokens, 1)
1044 else:
1045 avg_loss = total_loss_sum / max(total_tokens_local, 1)
1046 else:
1047 # rank_average: each rank averages its micro-batches, then
1048 # ``all_reduce`` averages across DP ranks. Divisor is the DP
1049 # group size (NOT global world_size); under TP/EP/PP/CP the
1050 # all-reduce only sums DP ranks, so dividing by world_size
1051 # would systematically under-report the loss.
1052 local_mean = total_loss_arith_sum / max(num_micro, 1)
1053 dp_size = self._dp_group_info.rank_size
1054 if dp_size > 1:
1055 ls = platform.full((1,), local_mean).to(self.device)
1056 platform.all_reduce(ls, self._dp_group_info)
1057 avg_loss = ls.item() / dp_size
1058 else:
1059 avg_loss = local_mean
1061 return {"loss": avg_loss, "grad_norm": grad_norm.item()}
1063 def train(self):
1064 """Main training loop: epoch → step → micro-batch.
1066 Dispatches callbacks at each lifecycle point (explicit mode).
1067 on_train_begin is called first — CheckpointCallback uses it to restore
1068 state.global_step from a saved checkpoint, so the loop below will
1069 correctly skip already-completed steps.
1070 """
1071 logger.info_rank0(
1072 "Training starts: max_steps=%d, epochs=%d",
1073 self.state.max_steps,
1074 getattr(self.args.train, 'num_train_epochs', 1),
1075 )
1076 # on_train_begin runs checkpoint resume — state.global_step may be
1077 # updated to the resumed step before the loop starts.
1078 self.on_train_begin()
1079 num_epochs = getattr(self.args.train, 'num_train_epochs', 1)
1081 if self.state.global_step > 0:
1082 logger.info_rank0(
1083 "Resuming training from step %d", self.state.global_step,
1084 )
1086 for epoch in range(num_epochs):
1087 if self.state.global_step >= self.state.max_steps:
1088 break
1089 self.state.epoch = epoch
1090 if self.sampler is not None:
1091 self.sampler.set_epoch(epoch)
1092 self.on_epoch_begin()
1094 # Build micro-batch iterator from the stateful dataloader.
1095 # StatefulDataLoader tracks iterator position internally,
1096 # so after resume it skips already-consumed batches.
1097 data_iterator = self._make_micro_batch_iterator()
1099 # Drive the loop on the live ``global_step`` so total training
1100 # never exceeds ``max_steps`` regardless of ``num_train_epochs``
1101 # or resume offset.
1102 while self.state.global_step < self.state.max_steps:
1103 self.on_step_begin()
1104 try:
1105 metrics = self.train_step(data_iterator)
1106 except StopIteration:
1107 logger.info_rank0("Epoch %d: dataloader exhausted", epoch)
1108 break
1110 self.on_step_end(
1111 loss=metrics["loss"],
1112 grad_norm=metrics["grad_norm"],
1113 )
1115 self.on_epoch_end()
1117 self.on_train_end()
1118 destroy_process_group()
1119 logger.info_rank0("Training completed")
1121 # ------------------------------------------------------------------
1122 # Helpers
1123 # ------------------------------------------------------------------
1125 def _make_micro_batch_iterator(self):
1126 """Yield lists of micro-batches from the stateful dataloader.
1128 Groups ``self._grad_accum`` consecutive batches into a list for
1129 gradient accumulation. The underlying ``StatefulDataLoader`` tracks
1130 iteration position, so checkpoint/resume skips consumed batches.
1131 """
1132 batch_buffer = []
1133 for batch in self.train_dataloader:
1134 batch_buffer.append(batch)
1135 if len(batch_buffer) >= self._grad_accum:
1136 yield batch_buffer
1137 batch_buffer = []
1138 if batch_buffer:
1139 yield batch_buffer
1141 def _get_layers(self) -> list:
1142 """Return the repeating layers for FSDP/AC wrapping.
1144 Default: ``model.layers`` (standard transformer convention).
1145 Override in subclass for models with different structure.
1146 """
1147 if hasattr(self.model, 'layers'):
1148 return list(self.model.layers)
1149 raise ValueError(
1150 f"Model {type(self.model).__name__} has no .layers attribute. "
1151 f"Either add self.layers to the model, or override _get_layers() "
1152 f"in the Trainer subclass."
1153 )
1155 def _get_combined_dp_group(self):
1156 """Return the combined data-parallel ProcessGroup for trainer all-reduce.
1158 Prefers the ``"loss"`` flatten alias registered by
1159 ``ParallelDims.build_mesh`` (folds CP into the DP group when CP is
1160 active so token-count denominators include CP-sharded contributions).
1161 Falls back to ``"dp"``, then to the legacy ``dp_shard`` /
1162 ``dp_replicate`` axes for callers that built a custom mesh.
1163 """
1164 for name in ("loss", "dp", "dp_shard", "dp_replicate"):
1165 try:
1166 return self.mesh.get_group(name)
1167 except (KeyError, ValueError):
1168 continue
1169 return self.mesh.get_group()
1171 def _build_fsdp_kwargs(self) -> dict:
1172 """Build kwargs for ``fully_shard`` calls (dense parameters).
1174 For expert parameters when EP > 1, use ``_build_expert_fsdp_kwargs``.
1175 """
1176 for name in ("dp_shard", "dp", "dp_replicate"):
1177 try:
1178 dp_mesh = self.mesh[name]
1179 break
1180 except (KeyError, TypeError):
1181 continue
1182 else:
1183 dp_mesh = self.mesh
1184 kwargs = {"mesh": dp_mesh}
1186 reshard = getattr(self.args.train.accelerator, 'reshard_after_forward', True)
1187 kwargs["reshard_after_forward"] = reshard
1189 return kwargs
1191 def _build_expert_fsdp_kwargs(self) -> dict:
1192 """Build kwargs for ``fully_shard`` calls on expert parameters.
1194 When EP > 1, expert parameters are sharded across the EP group
1195 with a separate mesh dimension. Falls back to dense FSDP kwargs
1196 if EP is not enabled.
1197 """
1198 if not self.parallel_dims.ep_enabled:
1199 return self._build_fsdp_kwargs()
1201 try:
1202 ep_mesh = self.mesh["ep"]
1203 except (KeyError, TypeError):
1204 logger.warning("EP=%d but no 'ep' dimension in mesh, falling back to dp mesh",
1205 self.parallel_dims.ep)
1206 return self._build_fsdp_kwargs()
1208 kwargs = {"mesh": ep_mesh}
1209 reshard = getattr(self.args.train.accelerator, 'reshard_after_forward', True)
1210 kwargs["reshard_after_forward"] = reshard
1211 return kwargs
1213 def _materialize_and_init_shards(self) -> None:
1214 """Materialize meta-device parameters/buffers to real device in-place.
1216 After ``fully_shard`` on a meta-device model, each rank's parameters
1217 are meta DTensor shards **and FSDP2 holds internal views into those
1218 meta storages** (flat_param / unsharded buffer). Replacing the
1219 ``DTensor._local_tensor`` attribute leaves FSDP's internal views
1220 pointing at the old meta storage, so the first forward's all-gather
1221 still hits meta → ``c10d::_allgather_base_`` raises.
1223 PyTorch's ``nn.Module.to_empty(device=...)`` is the FSDP2-safe path:
1224 it walks every parameter/buffer (including DTensor shards) and
1225 **allocates real device storage in-place via ``torch.empty_like``**,
1226 preserving every existing view. After ``to_empty``, storage is
1227 uninitialised — we init on the local shard with kaiming_uniform for
1228 weights, zero for biases / 1-D / buffers.
1230 Reference: PyTorch FSDP2 meta-init docs;
1231 ``trainer.py`` post-``fully_shard`` init pattern.
1232 """
1233 device_type = platform.device_type()
1234 # Step 1: meta → real storage, in-place (FSDP-views preserved).
1235 self.model.to_empty(device=device_type)
1236 self._materialize_replicate_params(device_type)
1237 # Step 2: init the local shard of every param (and zero every buffer).
1238 param_count = self._init_local_shards()
1239 # Re-derive buffers wiped by ``to_empty`` (e.g. ``inv_freq``);
1240 # without this RoPE silently returns identity rotation.
1241 for module in self.model.modules():
1242 if hasattr(module, "reset_inv_freq"):
1243 module.reset_inv_freq()
1244 # Re-tie weights — ``to_empty`` gives every nn.Parameter fresh
1245 # storage so ``__init__``-time ties are broken. Must happen before
1246 # ``lazy_init`` re-wraps params as DTensor (non-leaf), which would
1247 # cause ``register_parameter`` to reject the assignment.
1248 if hasattr(self.model, "tie_weights"):
1249 self.model.tie_weights()
1250 # ``to_empty`` strips DTensor; ``lazy_init`` re-wraps shards before
1251 # ``_load_weights`` / optimizer step see the params (the forward
1252 # pre-hook does the same later, but the loader needs DTensor first).
1253 reset_count = self._lazy_init_hsdp_modules()
1254 logger.info_rank0(
1255 "Meta → real on %s: to_empty + kaiming/zero init on %d params; "
1256 "FSDP lazy_init re-wrapped %d modules back to DTensor",
1257 device_type, param_count, reset_count,
1258 )
1260 def _iter_hsdp_states(self):
1261 """Yield the HSDP state attached to every HSDP-wrapped submodule."""
1262 for module in self.model.modules():
1263 if not isinstance(module, HSDPModule):
1264 continue
1265 scheduler = getattr(module, 'hsdp_scheduler', None)
1266 state = getattr(scheduler, 'hsdp_state', None) if scheduler else None
1267 if state is None:
1268 continue
1269 yield state
1271 def _materialize_replicate_params(self, device_type: str) -> None:
1272 """``to_empty`` skips ``replicate_params`` whose ``_local_tensor`` lives
1273 outside ``module._parameters`` — materialize them manually.
1274 """
1275 for state in self._iter_hsdp_states():
1276 for hsdp_param in getattr(state, 'replicate_params', []) or []:
1277 local = getattr(hsdp_param.sharded_param, "_local_tensor", None)
1278 if local is not None and local.is_meta:
1279 new_local = torch.empty_like(local, device=device_type)
1280 hsdp_param.sharded_param._local_tensor = new_local # pylint: disable=W0212
1282 def _init_local_shards(self) -> int:
1283 """Init local shard of every param (kaiming for >=2D, zero else); zero buffers."""
1284 param_count = 0
1285 with torch.no_grad():
1286 for _, param in self.model.named_parameters():
1287 local = param._local_tensor if hasattr(param, '_local_tensor') else param # pylint: disable=W0212
1288 if local.is_meta:
1289 continue
1290 if local.dim() >= 2:
1291 torch.nn.init.kaiming_uniform_(local)
1292 else:
1293 torch.nn.init.zeros_(local)
1294 param_count += 1
1295 for _, buf in self.model.named_buffers():
1296 if buf is not None:
1297 buf.zero_()
1298 return param_count
1300 def _lazy_init_hsdp_modules(self) -> int:
1301 """Re-wrap HSDP shards into DTensor so loader / optimizer see them."""
1302 reset_count = 0
1303 for state in self._iter_hsdp_states():
1304 if hasattr(state, 'lazy_init'):
1305 state.lazy_init()
1306 reset_count += 1
1307 return reset_count
1309 def _load_weights(self, weights_path: str) -> None:
1310 """Load pre-trained weights from ``weights_path`` into the (possibly sharded) model.
1312 Uses hyper's distributed checkpoint ``load`` API so that each rank only
1313 reads the shard it owns. Falls back to a plain ``torch.load`` + partial
1314 ``load_state_dict`` for single-file checkpoints (e.g. safetensors).
1316 Args:
1317 weights_path: Path to a directory containing a distributed checkpoint,
1318 or a single ``.pt`` / ``.bin`` file.
1319 """
1320 logger.info_rank0("Loading weights from %s", weights_path)
1321 try:
1322 if os.path.isdir(weights_path):
1323 hf_index = os.path.join(weights_path, "model.safetensors.index.json")
1324 # Delegate model-specific renaming / expert-splitting to
1325 # the per-spec ``state_dict_adapter``.
1326 adapter_cls = getattr(self.spec, "state_dict_adapter", None)
1327 if os.path.isfile(hf_index) and adapter_cls is not None:
1328 self._load_hf_safetensors(weights_path, adapter_cls)
1329 else:
1330 self._load_hyper_dcp(weights_path)
1331 else:
1332 self._load_single_file(weights_path)
1333 logger.info_rank0("Weights loaded from %s", weights_path)
1334 except Exception as exc:
1335 raise RuntimeError(
1336 f"Failed to load weights from {weights_path}: {exc}. "
1337 "weights_path was provided so silent random-init fallback is unsafe — "
1338 "uniform-logits loss would corrupt downstream training metrics."
1339 ) from exc
1341 def _load_hf_safetensors(self, weights_path: str, adapter_cls) -> None:
1342 """Load HF safetensors via spec's ``state_dict_adapter``; drop shape mismatches."""
1343 # Cast loaded params down to the checkpoint's advertised dtype so the
1344 # fp32 master matches what forward consumes.
1345 load_dtype = self._resolve_hf_load_dtype(weights_path)
1346 adapter = adapter_cls()
1347 hf_sd = adapter.load_hf_state_dict(
1348 weights_path, self.model.config, dtype=load_dtype,
1349 )
1350 valid_sd, dropped, missing, unexpected = self._validate_hf_state_dict(hf_sd)
1351 if dropped:
1352 logger.warning(
1353 "Dropped %d keys due to shape mismatch (first 5: %s)",
1354 len(dropped), dropped[:5],
1355 )
1356 # Derive missing/unexpected ourselves — ``HSDPModule.load_state_dict``
1357 # returns ``None``.
1358 self.model.load_state_dict(valid_sd, strict=False)
1359 model_name = getattr(self.args.model, "name", "")
1360 logger.info_rank0(
1361 "HF (%s) load: %d tensors into hyper model",
1362 model_name, len(valid_sd),
1363 )
1364 if missing:
1365 logger.warning(
1366 "Missing (randomly initialised): %d keys, e.g. %s ...",
1367 len(missing), missing[:5],
1368 )
1369 if unexpected:
1370 logger.warning(
1371 "Unexpected (ignored): %d keys, e.g. %s ...",
1372 len(unexpected), unexpected[:5],
1373 )
1375 def _resolve_hf_load_dtype(self, weights_path: str):
1376 """Resolve the dtype to cast loaded HF tensors to (matches checkpoint config)."""
1377 dtype_map = {
1378 'bfloat16': torch.bfloat16, 'bf16': torch.bfloat16,
1379 'float16': torch.float16, 'fp16': torch.float16,
1380 'float32': torch.float32, 'fp32': torch.float32,
1381 }
1382 cfg_dtype = (
1383 getattr(self.model.config, 'dtype', None)
1384 or getattr(self.model.config, 'torch_dtype', None)
1385 )
1386 if cfg_dtype is None:
1387 cfg_json = os.path.join(weights_path, 'config.json')
1388 if os.path.isfile(cfg_json):
1389 try:
1390 with open(cfg_json, 'r', encoding='utf-8') as f:
1391 cfg = json.load(f)
1392 cfg_dtype = cfg.get('dtype') or cfg.get('torch_dtype')
1393 except (OSError, json.JSONDecodeError):
1394 cfg_dtype = None
1395 if isinstance(cfg_dtype, str):
1396 return dtype_map.get(cfg_dtype)
1397 if isinstance(cfg_dtype, torch.dtype):
1398 return cfg_dtype
1399 return None
1401 def _validate_hf_state_dict(self, hf_sd: dict):
1402 """Strip wrapper segments and drop tensors whose shape differs from the model.
1404 Pre-validate shapes: ``load_state_dict`` aborts on the first mismatch
1405 and leaves later keys un-loaded.
1407 Returns:
1408 ``(valid_sd, dropped, missing, unexpected)``.
1409 """
1410 # Strip wrapper segments (e.g. ``_checkpoint_wrapped_module``) so
1411 # loader keys match ``named_parameters`` paths.
1412 wrapper_segments = ("._checkpoint_wrapped_module",)
1414 def _strip(k: str) -> str:
1415 for s in wrapper_segments:
1416 k = k.replace(s, "")
1417 return k
1418 logical_to_real = {}
1419 real_to_param = {}
1420 for name, param in self.model.named_parameters():
1421 logical_to_real[_strip(name)] = name
1422 real_to_param[name] = param
1423 valid_sd: dict = {}
1424 dropped: list = []
1425 for hf_name, hf_tensor in hf_sd.items():
1426 real_name = logical_to_real.get(hf_name)
1427 if real_name is None:
1428 continue
1429 param = real_to_param.get(real_name)
1430 if param is None:
1431 continue
1432 tgt = tuple(param.shape)
1433 src = tuple(hf_tensor.shape)
1434 if src == tgt:
1435 valid_sd[real_name] = hf_tensor
1436 else:
1437 dropped.append((real_name, src, tgt))
1438 param_names = set(real_to_param.keys())
1439 loaded_names = set(valid_sd.keys())
1440 missing = sorted(param_names - loaded_names)
1441 unexpected = sorted(loaded_names - param_names)
1442 return valid_sd, dropped, missing, unexpected
1444 def _load_hyper_dcp(self, weights_path: str) -> None:
1445 """Load weights from hyper's own DCP checkpoint format."""
1446 model_sd = self.model.state_dict()
1447 dcp_load(model_sd, checkpoint_id=weights_path, use_collectives=False)
1448 self.model.load_state_dict(model_sd)
1450 def _load_single_file(self, weights_path: str) -> None:
1451 """Load weights from a single ``.pt`` / ``.safetensors`` / ``.bin`` file."""
1452 sd = torch.load(weights_path, map_location="cpu", weights_only=True)
1453 missing, unexpected = self.model.load_state_dict(sd, strict=False)
1454 if missing:
1455 logger.warning("Missing keys when loading weights: %s", missing)
1456 if unexpected:
1457 logger.warning("Unexpected keys when loading weights: %s", unexpected)
1459 def _maybe_toggle_reshard(self, micro_step: int, num_micro_steps: int):
1460 """Toggle FSDP reshard_after_backward for gradient accumulation optimization.
1462 During gradient accumulation, skip resharding between micro-steps to avoid
1463 redundant all-gather. Only reshard after the last micro-step.
1464 """
1465 if not isinstance(self.model, HSDPModule) or num_micro_steps <= 1:
1466 return
1467 if micro_step == 0:
1468 self.model.set_reshard_after_backward(False)
1469 elif micro_step == num_micro_steps - 1:
1470 self.model.set_reshard_after_backward(True)