Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / base.py: 57%
487 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates
2# Copyright 2026 Huawei Technologies Co., Ltd
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
16"""
17Base Trainer class for distributed training.
19This module provides the BaseTrainer class which serves as the foundation
20for all trainer implementations. Subclasses can override specific methods
21to customize training behavior.
23Features:
24 - Callback system for extensible training hooks
25 - Distributed training support
26 - Gradient accumulation
27 - Checkpointing
28"""
29# pylint: disable=forbidden-backend-import
31from __future__ import annotations
33import json
34import logging
35import os
36from abc import ABC
37from collections import defaultdict
38from contextlib import nullcontext
39from functools import partial
40from typing import TYPE_CHECKING, Any, Dict, List, Optional
42import torch
43import torch.distributed as dist
44from torch.distributed.checkpoint.stateful import Stateful
45from torch.optim.lr_scheduler import LRScheduler
46from torch.utils.data import Dataset
47from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerBase
48from transformers.modeling_outputs import ModelOutput
50from hyper_parallel import HSDPModule, SkipDTensorDispatch
51from hyper_parallel.compile.compiler import GraphCompiler
52from hyper_parallel.core.tensor_parallel import loss_parallel
53from hyper_parallel.core.utils import clip_grad_norm_
54from hyper_parallel.trainer.config import (
55 OptimizerSwapConfig,
56 TrainerConfig,
57 normalize_distributed_setup_overrides,
58 save_configs,
59)
60from hyper_parallel.data.dataset_logging import enable_dataset_logging
61from hyper_parallel.data.batching import build_dataloader
62from hyper_parallel.trainer.runtime.distributed import (
63 get_global_rank_safe,
64 get_local_rank_safe,
65 get_world_size_safe,
66)
67from hyper_parallel.trainer.runtime.distributed import (
68 create_distributed_setup_from_config,
69 destroy_process_group,
70 initialize_distributed,
71)
72from hyper_parallel.trainer.runtime.logging import setup_logging
73from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
74from hyper_parallel.trainer.runtime.metrics import mean_global_loss
75from hyper_parallel.components.functional.aux_loss import aux_loss_scale_context, bind_aux_loss_scale
76from hyper_parallel.components.modules.moe_utils import MoEMonitorCallback
77from hyper_parallel.models._transformers.loss_parallel import causal_lm_loss_parallel
78from hyper_parallel.models._transformers.model_builder import _validate_optimize_dtype
79from hyper_parallel.models.external_state import get_model_external_state
80from hyper_parallel.components.losses.model_output import ModelOutputLoss
81from hyper_parallel.core.optimizer import (
82 ChainedOptimizer,
83 SwapOptimizerConfig,
84 swap_optimizer,
85)
86from hyper_parallel.components.optim.mixed_precision_optimizer import (
87 Float16OptimizerWithFloat16Params,
88 MixedPrecisionOptimizer,
89)
90from hyper_parallel.trainer.runtime import fsdp as fsdp_runtime
91from hyper_parallel.trainer.runtime import model_integration as model_integration_runtime
92from hyper_parallel.trainer.runtime.data_iterator import HyperIter
93from hyper_parallel.trainer.runtime.logging import enable_third_party_logging
94from hyper_parallel.trainer.runtime.memory import empty_cache, print_device_mem_info
95from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
96from hyper_parallel.trainer.runtime.random import enable_high_precision_for_bf16, set_seed
97from hyper_parallel.trainer.runtime.device import ( # pylint: disable=syntax-error
98 get_device_type,
99 get_torch_device,
100 synchronize,
101)
103from hyper_parallel.trainer.callbacks import (
104 EnvironMeterCallback,
105 EvaluateCallback,
106 GarbageCollectionCallback,
107 LoggingCallback,
108 ProfilerCallback,
109 TqdmCallback,
110 CheckpointerCallback,
111 TrainerState,
112)
114logger = logging.getLogger(__name__)
116if TYPE_CHECKING:
117 from hyper_parallel.data.text.chat_template import ChatTemplate
120def _to_swap_optimizer_config(config: OptimizerSwapConfig) -> SwapOptimizerConfig:
121 """Translate the trainer-facing swap section into the core swap config.
123 Args:
124 config: Resolved ``optimizer.swap`` configuration section.
126 Returns:
127 The backend-neutral swap configuration consumed by ``swap_optimizer``.
128 """
129 return SwapOptimizerConfig(
130 swap_times=config.swap_times,
131 state_keys=config.state_keys,
132 min_numel=config.min_numel,
133 # ``None`` means "leave it to the backend default", not "off": the core
134 # config declares ``packed_swap: bool = True``, so pass that default
135 # through rather than letting ``None`` read as a hard False.
136 packed_swap=True if config.packed_swap is None else config.packed_swap,
137 )
140def _attach_optimizer_swap(optimizer: Any, config: OptimizerSwapConfig) -> Any:
141 """Attach optimizer-state swap to every Adam/AdamW leaf of ``optimizer``.
143 Every YAML optimizer target returns a ``ChainedOptimizer`` of leaf
144 optimizers -- also for a single AdamW -- and the fp32 main-parameter wrapper
145 keeps the same leaf containers. The swap runtime wraps one concrete
146 Adam/AdamW and delegates everything else, so it is attached to the leaves
147 and the outer object keeps dispatching.
149 Args:
150 optimizer: Optimizer built by the configured target, optionally already
151 wrapped for fp32 main parameters.
152 config: Resolved ``optimizer.swap`` configuration section.
154 Returns:
155 The optimizer carrying swap-wrapped leaves. A leaf the swap runtime does
156 not support fails here, before the first training step.
157 """
158 swap_config = _to_swap_optimizer_config(config)
159 chained = optimizer.optimizer if isinstance(optimizer, MixedPrecisionOptimizer) else optimizer
160 if not isinstance(chained, ChainedOptimizer):
161 return swap_optimizer(optimizer, swap_config)
163 names = list(chained.optimizers_dict)
164 leaves = [swap_optimizer(chained.optimizers_dict[name], swap_config) for name in names]
165 # In place: the fp32 main-parameter wrapper aliases both containers.
166 chained.chained_optimizers[:] = leaves
167 for name, leaf in zip(names, leaves):
168 chained.optimizers_dict[name] = leaf
169 return optimizer
172class BaseTrainer(Stateful, ABC):
173 """
174 Base trainer class for distributed model training.
176 This class provides the core training infrastructure including:
177 - Distributed initialization and parallelism setup
178 - Model, optimizer, and scheduler initialization
179 - Training step execution with gradient accumulation
180 - Checkpointing and fault tolerance
181 - Metrics logging
183 Subclasses can override the following methods to customize behavior:
184 - `_build_distributed_setup()`: Connect AutoModel DistributedSetup/MeshContext
185 - `_build_model()`: Connect AutoModel from_pretrained/parallelize
186 - `_build_model_assets()`: Build tokenizer/processor assets
187 - `_build_data_transform()`: Build modality-specific sample transforms
188 - `forward_backward_step()`: Customize forward/backward logic
189 - `train_step()`: Customize training step execution
190 - `train()`: Train the model
192 Callback Hooks:
193 The trainer calls callback methods at various stages:
194 - evaluate_callback: evaluation callback
195 - trace_callback: tracing callback (meter, wandb, tqdm, profile)
196 - checkpoint_callback: checkpointing callback
197 """
199 # Core configs
200 config: TrainerConfig
201 device: torch.device
202 # Attached at runtime by the model-integration pipeline (the Omni dataset builder sets
203 # ``data_transform``, the integration adapter attaches ``get_batch`` and
204 # ``num_micro_batches``): bare annotations, which assign nothing, so runtime behaviour is
205 # unchanged -- they exist so the type checker sees the contract this class reads.
206 data_transform: Any
207 get_batch: Any
208 num_micro_batches: int
210 # AutoModel distributed setup.
211 # These objects have one owner: ``../build_options.py`` and
212 # ``../components/distributed/mesh.py``. Trainer code consumes them but must
213 # not create a second DeviceMesh or a VeOmni ParallelState.
214 dist_info: Any
215 distributed_setup: Any
216 mesh_context: Any
217 device_mesh: Any
218 moe_mesh: Any
220 # Data
221 train_dataset: Dataset
222 collate_fn: Any
223 train_dataloader: Any
225 # Model
226 model: PreTrainedModel = None
227 model_parts: Optional[List[torch.nn.Module]] = None
228 hsdp_model_parts: List[torch.nn.Module] = []
229 fsdp_runtime_diagnostics: Optional[fsdp_runtime.FSDPRuntimeDiagnostics] = None
230 model_integration: Optional[model_integration_runtime.ModelIntegrationSession] = None
231 model_config: PretrainedConfig = PretrainedConfig()
232 tokenizer: PreTrainedTokenizerBase = None
233 processor: Any = None
234 chat_template: Optional[ChatTemplate] = None
235 model_assets: List[Any] = []
237 # Training components
238 optimizer = None
239 lr_scheduler: LRScheduler = None
240 loss_fn: Optional[torch.nn.Module] = None
242 # Training context
243 model_fwd_context: Any
244 model_bwd_context: Any
246 # # Runtime metrics, controlled by trace_callback
247 # environ_meter: helper.EnvironMeter # see in trace_callback.EnvironMeterCallback
248 # step_env_metrics: Dict[str, Any] # mfu, flops, tokens, etc
249 # step_train_metrics: Dict[str, Any] # loss, grad_norm, lr, etc
251 # # Checkpointer
252 # checkpointer: CheckpointerBase # see in checkpoint_callback.CheckpointerCallback
254 # Callback system
255 state: TrainerState
257 # Default seed when training.seed is omitted.
258 default_seed: int = 42
260 def __init__(self, config: TrainerConfig) -> None:
261 """
262 Initialize the trainer.
264 Args:
265 args: Global Arguments
266 Should have attributes: model, data, train
267 model: ModelArguments
268 data: DataArguments
269 train: TrainingArguments
270 """
272 self.config: TrainerConfig = config
274 # General setup owns distributed initialization as its first stage.
275 # ``_build_distributed_setup`` remains a reserved backend hook; the
276 # trainer only defines its contract here.
277 self._setup()
278 memory_profiler.reset(
279 self.config.memory,
280 global_rank=self.global_rank,
281 tp_rank=self.mesh.tp_rank,
282 dp_rank=self.mesh.dp_rank,
283 )
285 # Reserved atomic model-build interface.
286 #
287 # Owner:
288 # - ``../_transformers/auto_model.py::HyperAutoModel*.from_pretrained``
289 # - ``../_transformers/model_builder.py`` for materialization,
290 # checkpoint loading, PEFT/quantization, and parallelize.
291 #
292 # Contract:
293 # - consume the ``DistributedSetup`` created above;
294 # - set ``model`` and the resolved checkpoint ``model_config``;
295 # - return a materialized, weight-loaded, already-parallelized model;
296 # - never call a second trainer-side ``build_parallelize_model``.
297 self._build_model()
298 self._build_loss()
300 # Build trainer-owned data components after the model finalizes parameters and sharding.
301 self._build_model_assets()
302 self._build_data_transform()
303 self._build_dataset()
304 self._build_collate_fn()
305 self._build_dataloader()
306 self._compute_train_iters()
308 self._build_optimizer()
309 self.model_integration.attach_optimizer(self.optimizer)
310 self._build_lr_scheduler()
311 self._build_training_context()
312 self._init_callbacks()
314 def _setup(self):
315 """Initialize logging, distributed state, and the local device."""
316 # log args
317 setup_logging()
319 # Initialize distributed process group
320 backend = getattr(self.config.training, "backend", "nccl")
321 if backend == "hccl" and self.config.training.enable_full_determinism:
322 # HCCL reads this when creating communicators, before set_seed runs below.
323 os.environ["HCCL_DETERMINISTIC"] = "true"
324 initialize_distributed(backend=backend)
326 self.local_rank = get_local_rank_safe()
327 self.global_rank = get_global_rank_safe()
328 self.world_size = get_world_size_safe()
329 if self.config.debug.check_dataset:
330 enable_dataset_logging(self.config.debug.check_dataset)
332 if self.config.training.seed is None:
333 self.config.training.seed = self.default_seed
335 if self.global_rank == 0:
336 logger.info(
337 "Trainer config:\n%s",
338 json.dumps(self.config.to_dict(), indent=2),
339 )
341 # init distributed environment
342 device_str = f"{get_device_type()}:{self.local_rank}"
343 get_torch_device().set_device(device_str)
344 self.device = torch.device(device_str)
346 logger.info("Process rank: %s, world size: %s", self.global_rank, self.world_size)
348 # Register distributed_setup; the Trainer-side desugar fills its
349 # normalized plan_overrides/module_replacements (05 §15.2.6) so the
350 # AutoModels build pipeline never sees raw YAML PlanOverride DTOs.
351 self.distributed_setup = create_distributed_setup_from_config(self.config)
352 normalize_distributed_setup_overrides(self.distributed_setup, self.config)
353 self.mesh = self.distributed_setup.mesh_context
354 self.device_mesh = self.mesh.device_mesh
355 self.dp_cp_mesh = self.mesh.dp_cp_mesh
357 # Set random seed
358 set_seed(self.config.training.seed, self.config.training.enable_full_determinism)
360 # Enable high precision for bf16
361 enable_high_precision_for_bf16()
363 # Enable third party logging
364 if self.local_rank == 0:
365 enable_third_party_logging()
367 # Save arguments
368 if self.global_rank == 0:
369 save_configs(self.config, self.config.checkpoint.checkpoint_dir)
371 def _build_model(self) -> None:
372 """Build the model and derive Trainer-owned runtime state."""
373 self.peft_config = self.config.peft
374 self.model = self.config.model.build(
375 distributed_setup=self.distributed_setup,
376 peft_config=self.peft_config,
377 activation_checkpoint=self.config.activation_checkpoint.mode,
378 activation_checkpoint_selection=self.config.activation_checkpoint.selection,
379 swap_inputs=getattr(self.config.activation_checkpoint, "swap_inputs", False),
380 allow_cache_entry_mutation=getattr(
381 self.config.activation_checkpoint, "allow_cache_entry_mutation", False),
382 activation_swap=self.config.activation_swap,
383 compile_config=self.config.compile,
384 codegen=self.config.codegen,
385 # The final dtype is applied inside the atomic build (05 stage-5
386 # item 5); the Trainer no longer patches it afterwards.
387 model_init_dtype=self.config.model_init_dtype,
388 codegen_config=self.config,
389 )
390 _validate_optimize_dtype(
391 self.model,
392 fp32_main_params=self.config.optimizer.fp32_main_params,
393 )
394 self.model_config = self.model.config
395 if self.global_rank == 0:
396 logger.info(
397 "Model config:\n%s",
398 json.dumps(self.model.config.to_dict(), indent=2),
399 )
400 if self.config.accelerator.loss_parallel:
401 self.model.loss_function = causal_lm_loss_parallel
402 model_parts = getattr(self.model, "parts", None)
403 self.model_parts = list(model_parts) if model_parts is not None else [self.model]
404 self.hsdp_model_parts = [
405 model_part
406 for model_part in self.model_parts
407 if isinstance(model_part, HSDPModule)
408 ]
409 self.model_integration = model_integration_runtime.build_model_integration_session(
410 self.model,
411 self.config,
412 self.global_rank,
413 self.mesh,
414 )
415 self.fsdp_runtime_diagnostics = (
416 fsdp_runtime.build_fsdp_runtime_diagnostics(self.hsdp_model_parts)
417 if getattr(self.config.debug, "check_fsdp_runtime", False)
418 else None
419 )
421 def _build_loss(self) -> None:
422 """Build the configured loss module or use the model-output default."""
423 loss_fn = (
424 self.config.loss_fn.build()
425 if self.config.loss_fn is not None
426 else ModelOutputLoss()
427 )
428 if not isinstance(loss_fn, torch.nn.Module):
429 raise ValueError("config.loss_fn must build a torch.nn.Module")
430 self.loss_fn = loss_fn
431 # Model-integrated objectives may need to intercept the model before its
432 # ordinary forward materializes an otherwise avoidable output tensor.
433 loss_model_binding_hook = getattr(self.loss_fn, "bind_model", None)
434 if callable(loss_model_binding_hook):
435 loss_model_binding_hook(self.model, self.distributed_setup)
437 def _build_model_assets(self) -> None:
438 """Require a concrete Trainer to build modality-specific model assets."""
439 raise NotImplementedError("Concrete Trainer must implement _build_model_assets")
441 def _build_data_transform(self) -> None:
442 """Require a concrete Trainer to build its sample transform."""
443 raise NotImplementedError("Concrete Trainer must implement _build_data_transform")
445 def _build_dataset(self) -> None:
446 """Build and assign one Dataset or a train-validation-test tuple."""
447 dataset_result = self.config.dataset.build(
448 transform=self.data_transform,
449 tokenizer=self.tokenizer,
450 mesh_context=self.mesh,
451 training_config=self.config.training,
452 )
454 if isinstance(dataset_result, tuple) and len(dataset_result) == 3:
455 self.train_dataset, self.valid_dataset, self.test_dataset = dataset_result
456 return
458 self.train_dataset = dataset_result
459 self.valid_dataset = self.test_dataset = None
461 def _build_collate_fn(self) -> None:
462 """Require a concrete Trainer to build its micro-batch collator."""
463 raise NotImplementedError("Concrete Trainer must implement _build_collate_fn")
465 def _build_dataloader(self) -> None:
466 """Build and assign train, validation, and test dataloaders."""
467 split_names = ("train", "valid", "test")
468 datasets = tuple(getattr(self, f"{split_name}_dataset", None) for split_name in split_names)
469 dataloaders, batch_samplers = build_dataloader(
470 self.config.dataloader,
471 datasets=datasets,
472 collate_fn=self.collate_fn,
473 mesh_context=self.mesh,
474 training_config=self.config.training,
475 max_seq_len=getattr(self.data_transform, "max_seq_len", None),
476 default_seed=self.default_seed,
477 )
478 for split_name, dataloader, batch_sampler in zip(split_names, dataloaders, batch_samplers):
479 setattr(self, f"{split_name}_dataloader", dataloader)
480 setattr(self, f"{split_name}_batch_sampler", batch_sampler)
482 def attach_model_integration_data_pipeline(self) -> None:
483 """Attach the final data pipeline through one shared Trainer path."""
484 self.model_integration.attach_data_pipeline(
485 runtime_adapter=getattr(self.get_batch, "runtime_input_adapter", None),
486 get_batch=self.get_batch,
487 model_assets=tuple(self.model_assets),
488 )
490 def _compute_train_iters(self) -> None:
491 """Resolve the run limit and the optimizer steps available per Dataset epoch."""
492 training_config = self.config.training
493 if training_config.train_iters is not None and training_config.train_samples is not None:
494 raise ValueError("configure only one of training.train_iters and training.train_samples")
496 if training_config.train_iters is None and training_config.train_samples is None:
497 raise ValueError("configure one of training.train_iters or training.train_samples")
499 train_iters = training_config.train_iters
500 if training_config.train_samples is not None:
501 train_iters = training_config.train_samples // training_config.global_batch_size
503 if train_iters <= 0:
504 raise ValueError("training configuration must produce at least one train iteration")
506 try:
507 micro_batches_per_epoch = len(self.train_dataloader)
508 except TypeError:
509 micro_batches_per_epoch = 0
511 if micro_batches_per_epoch:
512 train_steps = micro_batches_per_epoch // self.num_micro_batches
513 if train_steps <= 0:
514 raise ValueError("train DataLoader must provide one complete optimizer step per epoch")
515 else:
516 # Streaming and dynamic DataLoaders define an epoch only when their
517 # source iterator is exhausted.
518 train_steps = train_iters
520 self.train_iters = train_iters
521 self.train_steps = train_steps
522 self.train_epochs = (self.train_iters + self.train_steps - 1) // self.train_steps
524 def _build_optimizer(self) -> None:
525 """Build the configured optimizer with the runtime model context."""
526 config: TrainerConfig = self.config
527 external = get_model_external_state(self.model)
528 if external is not None:
529 def wrap_dense(dense_raw: Any) -> Any:
530 """Wrap dense parameters in FP32 main and optional state swap.
532 Args:
533 dense_raw: Constructed dense optimizer.
534 """
535 dense = (
536 Float16OptimizerWithFloat16Params(dense_raw, self.model)
537 if config.optimizer.fp32_main_params else dense_raw
538 )
539 if config.optimizer.swap.enabled:
540 dense = _attach_optimizer_swap(dense, config.optimizer.swap)
541 return dense
543 self.optimizer = external.build_optimizer(config.optimizer, wrap_dense=wrap_dense)
544 return
545 optimizer = config.optimizer.target.build(model=self.model).get_optimizer()
546 self.optimizer = (
547 Float16OptimizerWithFloat16Params(optimizer, self.model)
548 if config.optimizer.fp32_main_params
549 else optimizer
550 )
551 # Swap is attached after the fp32 main-parameter wrap so leaf optimizers
552 # already reference their final parameters, and the lr scheduler built
553 # right after reads the wrapped ``self.optimizer.param_groups``.
554 if config.optimizer.swap.enabled:
555 self.optimizer = _attach_optimizer_swap(self.optimizer, config.optimizer.swap)
557 def _build_lr_scheduler(self):
558 config: TrainerConfig = self.config
559 scheduler_iters = config.training.lr_scheduler_iters or self.train_iters
560 if scheduler_iters <= 0:
561 raise ValueError("training.lr_scheduler_iters must be positive when configured")
562 lr_scheduler = config.lr_scheduler.build(
563 optimizer=self.optimizer,
564 train_iters=scheduler_iters,
565 )
566 self.lr_scheduler = lr_scheduler.get_lr_scheduler()
568 def _build_training_context(self):
569 """Build training context for distributed training."""
570 if self.config.accelerator.loss_parallel:
571 tp_mesh = self.device_mesh["tp"]
572 self.model_fwd_context = partial(loss_parallel, mesh=tp_mesh)
573 else:
574 self.model_fwd_context = nullcontext()
575 self.model_bwd_context = nullcontext()
576 self.graph_compiler: Optional[GraphCompiler] = None
578 def _build_graph_compiler(self) -> None:
579 """Build the optional joint-graph compiler for this trainer runtime."""
580 if not self.config.compile.selects_graph_compiler():
581 return
582 self.graph_compiler = GraphCompiler(
583 model=self.model,
584 train_fn=self._graph_train_fn,
585 trainer_config=self.config,
586 device=self.device,
587 mesh_context=self.mesh,
588 )
590 def _graph_train_fn(
591 self,
592 model: torch.nn.Module,
593 model_inputs: Dict[str, Any],
594 labels: Optional[torch.Tensor],
595 current_token_counts: Dict[str, torch.Tensor],
596 step_token_counts: Dict[str, torch.Tensor],
597 ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
598 """Capture forward/loss with explicit runtime token counts."""
599 model_fwd_context = (
600 self.model_fwd_context()
601 if callable(self.model_fwd_context)
602 else self.model_fwd_context
603 )
604 with model_fwd_context:
605 outputs = model(**model_inputs, use_cache=False)
606 # Counts must be graph inputs: packed lengths and ignored-label counts
607 # can vary even when tensor shapes happen to match the first batch.
608 loss, loss_dict = self.postforward(
609 outputs, labels,
610 current_token_counts=current_token_counts,
611 step_token_counts=step_token_counts,
612 )
613 return loss, loss_dict
615 def _init_callbacks(self):
616 """Initialize callbacks."""
617 dp_cp_mesh = self.mesh.dp_cp_mesh
618 self.moe_monitor = MoEMonitorCallback(
619 self.model,
620 dp_group=dp_cp_mesh.get_group() if dp_cp_mesh is not None else None,
621 tp_group=self.device_mesh["tp"].get_group() if self.mesh.sequence_parallel else None,
622 )
623 self.state = TrainerState()
624 self.environ_meter_callback = EnvironMeterCallback(self)
625 self.tqdm_callback = TqdmCallback(self)
626 self.logging_callback = LoggingCallback(self)
627 self.evaluate_callback = EvaluateCallback(self)
628 self.garbage_collection_callback = GarbageCollectionCallback(self)
629 self.profiler_callback = ProfilerCallback(self)
630 self._callbacks = [
631 self.environ_meter_callback,
632 self.logging_callback,
633 self.tqdm_callback,
634 self.evaluate_callback,
635 self.garbage_collection_callback,
636 self.profiler_callback,
637 ]
638 # Registered to save, to restore, or both --- the two are independent, so
639 # a run that only loads an existing checkpoint still needs the callback.
640 if self.config.checkpoint.save_ckpt or self.config.checkpoint.restore_from is not None:
641 self._callbacks.append(CheckpointerCallback(self))
642 else:
643 logger.info(
644 "Checkpointing is inactive (save_ckpt=false, restore_from=None); "
645 "no checkpoint callback registered."
646 )
648 def on_train_begin(self) -> None:
649 """Run all registered callbacks at the start of training."""
650 for callback in self._callbacks:
651 callback.on_train_begin(self.state)
653 def on_train_end(self) -> None:
654 """Run all registered callbacks at the end of training."""
655 for callback in self._callbacks:
656 callback.on_train_end(self.state)
658 def on_epoch_begin(self) -> None:
659 """Run all registered callbacks at the start of an epoch."""
660 for callback in self._callbacks:
661 callback.on_epoch_begin(self.state)
663 def on_epoch_end(self) -> None:
664 """Run all registered callbacks at the end of an epoch."""
665 for callback in self._callbacks:
666 callback.on_epoch_end(self.state)
668 def on_step_begin(self, **kwargs: Any) -> None:
669 """Run all registered callbacks at the start of a training step."""
670 for callback in self._callbacks:
671 callback.on_step_begin(self.state, **kwargs)
673 def on_micro_step_begin(self, micro_batch: Dict[str, Any], **kwargs: Any) -> None:
674 """Run all registered callbacks before one forward-backward micro step.
676 Args:
677 micro_batch: Prepared inputs for the current micro step.
678 **kwargs: Additional callback context.
679 """
680 for callback in self._callbacks:
681 callback.on_micro_step_begin(self.state, micro_batch, **kwargs)
683 def on_step_end(
684 self,
685 loss: Optional[float] = None,
686 loss_dict: Optional[Dict[str, Any]] = None,
687 grad_norm: Optional[float] = None,
688 ) -> None:
689 """Run all registered callbacks at the end of a training step.
691 Args:
692 loss: Aggregated loss for the optimizer step.
693 loss_dict: Named loss values for the optimizer step.
694 grad_norm: Gradient norm measured before the optimizer update.
695 """
696 for callback in self._callbacks:
697 callback.on_step_end(self.state, loss=loss, loss_dict=loss_dict, grad_norm=grad_norm)
699 def preforward(self, micro_batch: List[Dict[str, Any]]) -> Dict[str, Any]:
700 """Preprocess micro batches before forward pass.
702 Tensors are moved to ``self.device`` non-blockingly. Nested dicts
703 (e.g. ``multimodal_metadata`` emitted by ``PackingCollator``) are
704 recursed so inner tensor values land on the device too; Python ints
705 / lists / etc. pass through unchanged.
707 Args:
708 micro_batch: Batch fields to move to the training device.
710 Returns:
711 Batch fields with tensors moved to the training device.
712 """
714 def _to_device(v: Any) -> Any:
715 if isinstance(v, torch.Tensor):
716 return v.to(self.device, non_blocking=True)
717 if isinstance(v, dict):
718 return {k: _to_device(vv) for k, vv in v.items()}
719 return v
721 # chunk_mbs_config = getattr(self.config.training, "chunk_mbs_config", None)
722 # self._chunk_mbs_ranges = build_chunk_mbs_ranges(micro_batch, chunk_mbs_config)
723 micro_batch = {k: _to_device(v) for k, v in micro_batch.items()} # type: ignore
724 # if getattr(self, "LOG_SAMPLE", True):
725 # self.LOG_SAMPLE = False
726 return micro_batch
728 def postforward(
729 self, outputs: ModelOutput, labels: Optional[torch.Tensor],
730 *,
731 current_token_counts: Dict[str, torch.Tensor],
732 step_token_counts: Dict[str, torch.Tensor],
733 ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
734 """Postprocess model outputs after forward pass.
736 Args:
737 outputs: Model outputs consumed by the configured loss.
738 labels: Labels associated with the current micro-batch.
739 current_token_counts: Per-loss token counts for the current micro-batch.
740 step_token_counts: Per-loss token counts for the optimizer step.
742 Returns:
743 Backward loss and named globally aggregated loss values.
744 """
745 local_loss = self.loss_fn(model_output=outputs, labels=labels)
746 foundation_loss = local_loss.get("foundation_loss") if isinstance(local_loss, dict) else local_loss
747 if foundation_loss is not None:
748 bind_aux_loss_scale(foundation_loss)
749 loss_dict: Dict[str, torch.Tensor] = mean_global_loss(
750 local_loss,
751 current_token_counts,
752 step_token_counts,
753 device_mesh=self.mesh,
754 )
755 loss = torch.stack(list(loss_dict.values())).sum()
756 # Model-declared extra metrics join the logged mapping only after
757 # the backward loss is fixed: log-only, never part of the loss.
758 extra_metrics = getattr(outputs, "extra_metrics", None)
759 if extra_metrics is None:
760 model_ref = getattr(self, "model", None)
761 extra_metrics = getattr(model_ref, "pending_extra_metrics", None)
762 if extra_metrics is not None:
763 model_ref.pending_extra_metrics = None
764 if extra_metrics:
765 for name, value in extra_metrics.items():
766 loss_dict[name] = value.detach().float()
767 return loss, loss_dict
769 def forward_backward_step(
770 self,
771 micro_batch: dict[str, torch.Tensor],
772 loss_inputs: Optional[dict[str, torch.Tensor]] = None,
773 ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
774 """Run forward and backward computation for one micro batch.
776 Args:
777 micro_batch: Inputs forwarded to the model.
778 loss_inputs: Loss-only labels and masks from the get-batch runtime. If
779 omitted, the legacy single-dictionary path is preserved.
781 Returns:
782 Backward loss and the named globally aggregated loss mapping.
783 """
784 channel_loss_callback = getattr(self, "channel_loss_callback", None)
785 micro_step_context = (
786 channel_loss_callback.micro_step_context(self.state, micro_batch)
787 if channel_loss_callback is not None
788 else nullcontext()
789 )
790 with micro_step_context, aux_loss_scale_context():
791 micro_batch = self.preforward(micro_batch)
792 # TextTrainer passes loss-only fields separately. Keep the legacy
793 # BaseTrainer/VLM one-dictionary call contract working as well.
794 if loss_inputs is None:
795 loss_inputs = {
796 name: micro_batch[name]
797 for name in ("labels", "shift_labels", "loss_mask")
798 if name in micro_batch
799 }
800 else:
801 loss_inputs = self.preforward(loss_inputs)
802 labels = loss_inputs.get("labels", micro_batch.get("labels"))
804 # A model-integrated loss can translate the public batch fields into
805 # model-family protocol arguments before the model forward starts.
806 prepare_model_inputs = getattr(
807 self.loss_fn,
808 "prepare_model_inputs",
809 None,
810 )
811 if callable(prepare_model_inputs):
812 micro_batch = prepare_model_inputs(micro_batch, loss_inputs)
813 if channel_loss_callback is not None:
814 channel_loss_callback.strip_model_inputs(micro_batch)
815 self.model_integration.record_batch(micro_batch, loss_inputs)
817 if self.graph_compiler is not None:
818 if self.config.training.empty_cache_before_backward:
819 logger.warning(
820 "training.empty_cache_before_backward is ignored in graph compiler mode "
821 "because forward and backward execute in one joint graph"
822 )
823 # The traced ``loss_dict`` values flow back as extra graph
824 # outputs (emitted after the gradients), so the logging
825 # metrics keep per-key parity with the eager path.
826 loss, loss_dict = self.graph_compiler.forward_backward(
827 model_inputs=micro_batch,
828 labels=labels,
829 current_token_counts=self.current_token_counts,
830 step_token_counts=self.step_token_counts,
831 )
832 else:
833 model_fwd_context = (
834 self.model_fwd_context()
835 if callable(self.model_fwd_context)
836 else self.model_fwd_context
837 )
838 with model_fwd_context:
839 outputs: ModelOutput = self.model(**micro_batch, use_cache=False)
841 # with use_parallel_state("base"):
842 loss, loss_dict = self.postforward(
843 outputs, labels,
844 current_token_counts=self.current_token_counts,
845 step_token_counts=self.step_token_counts,
846 )
847 # The loss graph owns everything required for backward. Releasing
848 # the model output here avoids retaining large vocabulary logits
849 # until the whole backward pass finishes.
850 del outputs
851 if self.config.training.empty_cache_before_backward:
852 empty_cache()
854 # Backward pass
855 with self.model_bwd_context:
856 loss.backward()
858 del micro_batch
859 return loss, loss_dict
861 def model_reshard(self, micro_step: int, num_micro_steps: int) -> None:
862 """Reshard model after backward pass; policy lives in ``runtime/fsdp.py``.
864 Args:
865 micro_step: Zero-based micro-step index.
866 num_micro_steps: Number of micro-steps in the optimizer step.
867 """
868 fsdp_runtime.model_reshard(self.hsdp_model_parts, self.config.fsdp_config, micro_step, num_micro_steps)
870 def _configure_fsdp_gradient_sync(self, micro_step: int, num_micro_steps: int):
871 """Configure FSDP gradient synchronization; policy lives in ``runtime/fsdp.py``."""
872 fsdp_runtime.configure_fsdp_gradient_sync(
873 self.hsdp_model_parts,
874 self.config.fsdp_config,
875 self.mesh.dp_replicate_size,
876 micro_step,
877 num_micro_steps,
878 )
880 def configure_fsdp_gradient_sync(self, micro_step: int, num_micro_steps: int) -> None:
881 """Configure FSDP gradient synchronization for an external training loop.
883 Args:
884 micro_step: Zero-based micro-step index.
885 num_micro_steps: Number of micro-steps in the optimizer step.
886 """
887 self._configure_fsdp_gradient_sync(micro_step, num_micro_steps)
889 def begin_fsdp_runtime_diagnostics(self, micro_step: int) -> None:
890 """Begin observing one FSDP forward/backward micro-step when enabled.
892 Args:
893 micro_step: Microbatch index.
894 """
895 if self.fsdp_runtime_diagnostics is not None:
896 self.fsdp_runtime_diagnostics.begin_micro_step(
897 self.state.global_step,
898 micro_step,
899 )
901 def validate_fsdp_runtime_before_optimizer(self) -> None:
902 """Fail before optimizer updates when FSDP prefetch state is unsafe."""
903 if self.fsdp_runtime_diagnostics is not None:
904 self.model_integration.record_fsdp_trace(
905 self.fsdp_runtime_diagnostics.trace_evidence()
906 )
907 self.fsdp_runtime_diagnostics.validate_before_optimizer_step()
909 def prepare_optimizer_step(self) -> Any:
910 """Run pre-optimizer diagnostics and clipping in their required order."""
911 self.model_integration.after_backward_before_clip()
912 self.validate_fsdp_runtime_before_optimizer()
913 max_grad_norm = self.config.training.max_grad_norm
914 external = get_model_external_state(self.model)
915 if external is not None:
916 preparation = external.prepare_optimizer_step(max_norm=max_grad_norm)
917 self.model_integration.after_clip(preparation.post_clip_norm)
918 return preparation.global_norm
919 grad_norm: Any = 0.0
920 post_clip_norm = None
921 if max_grad_norm > 0:
922 grad_norm = clip_grad_norm_(self.model, max_grad_norm)
923 if self.model_integration.runtime_enabled:
924 grad_norm_value = (
925 float(grad_norm.item())
926 if isinstance(grad_norm, torch.Tensor)
927 else float(grad_norm)
928 )
929 clip_coefficient = min(
930 float(max_grad_norm) / (grad_norm_value + 1.0e-6),
931 1.0,
932 )
933 post_clip_norm = grad_norm_value * clip_coefficient
934 elif self.model_integration.runtime_enabled:
935 grad_norm = clip_grad_norm_(self.model, float("inf"))
936 post_clip_norm = grad_norm
937 self.model_integration.after_clip(post_clip_norm)
938 return grad_norm
940 def step_optimizers_and_schedulers(self) -> None:
941 """Step every optimizer and scheduler, then expose the final state."""
942 optimizers = self.optimizer if isinstance(self.optimizer, list) else [self.optimizer]
943 for optimizer in optimizers:
944 # Keep zeros_like dispatched so lazily created optimizer states retain DTensor metadata.
945 with SkipDTensorDispatch(no_skip={torch.zeros_like}):
946 optimizer.step()
947 optimizer.zero_grad()
948 moe_monitor = getattr(self, "moe_monitor", None)
949 if moe_monitor is not None:
950 moe_monitor.on_step_end()
951 self.model_integration.after_optimizer()
953 schedulers = (
954 self.lr_scheduler
955 if isinstance(self.lr_scheduler, list)
956 else ([self.lr_scheduler] if self.lr_scheduler is not None else [])
957 )
958 for scheduler in schedulers:
959 scheduler.step()
961 def _end_model_integration_step(self, metrics: Dict[str, Any]) -> None:
962 """Persist structured step metrics and global input identity."""
963 if not self.model_integration.runtime_enabled:
964 return
965 resolved_metrics = dict(metrics)
966 optimizers = self.optimizer if isinstance(self.optimizer, list) else [self.optimizer]
967 learning_rates = []
968 for optimizer in optimizers:
969 inner = getattr(optimizer, "optimizer", optimizer)
970 leaves = getattr(inner, "chained_optimizers", (inner,))
971 for leaf in leaves:
972 learning_rates.extend(
973 group.get("lr")
974 for group in getattr(leaf, "param_groups", ())
975 if group.get("lr") is not None
976 )
977 if learning_rates:
978 resolved_metrics["lr"] = float(learning_rates[0])
979 token_counts = getattr(self, "step_token_counts", {})
980 token_values = []
981 for value in token_counts.values():
982 token_values.append(float(value.item() if torch.is_tensor(value) else value))
983 if token_values:
984 resolved_metrics["tokens"] = sum(token_values)
985 resolved_metrics["samples"] = self.config.training.global_batch_size
986 self.model_integration.end_step(resolved_metrics)
988 def train_step(
989 self,
990 data_iterator: Any,
991 ) -> Dict[str, float]:
992 """Execute one optimizer update from the next dataloader batch.
994 Args:
995 data_iterator: Iterator providing the next micro-batches.
997 Returns:
998 Aggregated loss and gradient norm.
999 """
1000 memory_profiler.step()
1001 micro_batches: List[Dict[str, Any]] = next(data_iterator)
1002 self.state.global_step += 1
1004 self.on_step_begin(micro_batches=micro_batches)
1005 self.model_integration.begin_step(self.state.global_step)
1007 # Forward and backward for each micro batch
1008 synchronize()
1010 total_loss = 0.0
1011 total_loss_dict = defaultdict(int)
1013 # token num for fixed_ce_loss in postforward
1014 self.step_token_counts = count_loss_token(micro_batches)
1015 num_micro_steps = len(micro_batches)
1016 # forward and backward pass with gradient_accumulationsteps
1017 for micro_step, micro_batch in enumerate(micro_batches):
1018 self.model_reshard(micro_step, num_micro_steps)
1019 self._configure_fsdp_gradient_sync(micro_step, num_micro_steps)
1020 self.begin_fsdp_runtime_diagnostics(micro_step)
1021 loss: torch.Tensor
1022 loss_dict: Dict[str, torch.Tensor]
1023 # token num for fixed_ce_loss in postforward
1024 self.current_token_counts = count_loss_token(micro_batch)
1025 loss, loss_dict = self.forward_backward_step(micro_batch)
1027 total_loss += loss.item()
1028 for k, v in loss_dict.items():
1029 total_loss_dict[k] += v.item()
1031 grad_norm = self.prepare_optimizer_step()
1032 self.step_optimizers_and_schedulers()
1034 grad_norm_value = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else float(grad_norm)
1035 self._end_model_integration_step({"loss": total_loss, "grad_norm": grad_norm_value})
1036 self.on_step_end(loss=total_loss, loss_dict=total_loss_dict, grad_norm=grad_norm_value)
1037 return {
1038 "loss": total_loss,
1039 "grad_norm": grad_norm_value,
1040 }
1042 def abort_profilers(self) -> None:
1043 """Release process-local profiling resources while preserving a training error."""
1044 try:
1045 memory_profiler.abort()
1046 # Cleanup must also preserve the training error if interrupted.
1047 except BaseException: # pylint: disable=broad-exception-caught
1048 logger.exception("Failed to stop memory profiling while aborting training.")
1049 profiler_callback = getattr(self, "profiler_callback", None)
1050 if profiler_callback is not None:
1051 try:
1052 profiler_callback.close()
1053 # Cleanup must also preserve the training error if interrupted.
1054 except BaseException: # pylint: disable=broad-exception-caught
1055 logger.exception("Failed to stop performance profiling while aborting training.")
1057 def destroy_distributed(self) -> None:
1058 """Synchronize all ranks and tear down the distributed process group."""
1059 if not dist.is_available() or not dist.is_initialized():
1060 return
1062 empty_cache()
1063 dist.barrier()
1065 synchronize()
1066 destroy_process_group()
1068 def train(self) -> None:
1069 """Run the configured training loop."""
1070 try:
1071 config: TrainerConfig = self.config
1072 self.on_train_begin()
1073 self.data_iterator = HyperIter(
1074 self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
1075 )
1076 logger.info(
1077 "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
1078 self.local_rank, self.state.global_step, self.train_iters, self.state.epoch, self.train_epochs,
1079 )
1081 try:
1082 start_epoch = self.state.epoch
1083 for epoch in range(start_epoch, self.train_epochs):
1084 if epoch != start_epoch:
1085 self.train_dataloader.set_epoch(epoch)
1086 self.data_iterator = HyperIter(
1087 self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
1088 )
1089 self.state.epoch = epoch
1091 self.on_epoch_begin()
1093 start_step = self.state.global_step - epoch * self.train_steps
1094 train_steps = min(self.train_steps, self.train_iters - epoch * self.train_steps)
1095 for _ in range(start_step, train_steps):
1096 try:
1097 self.train_step(self.data_iterator)
1098 except StopIteration:
1099 logger.info(
1100 "epoch:%s Dataloader finished with drop_last %s",
1101 epoch,
1102 config.dataloader.drop_last,
1103 )
1104 break
1106 self.on_epoch_end()
1107 self.state.epoch = epoch + 1
1109 print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")
1111 if config.dataloader.use_background_prefetcher:
1112 self.data_iterator.stop()
1114 memory_profiler.stop()
1115 self.on_train_end()
1116 finally:
1117 if config.dataloader.use_background_prefetcher:
1118 self.data_iterator.stop()
1120 except BaseException:
1121 try:
1122 memory_profiler.stop()
1123 # Snapshot export must preserve the training error even if interrupted.
1124 except BaseException: # pylint: disable=broad-exception-caught
1125 logger.exception("Failed to save memory snapshot after training error.")
1126 finally:
1127 self.abort_profilers()
1128 raise
1130 synchronize()
1132 self.destroy_distributed()