Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / base.py: 63%
556 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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 collections.abc import Callable
39from contextlib import nullcontext
40from functools import partial
41from typing import TYPE_CHECKING, Any, Dict, List, Optional
43import torch
44import torch.distributed as dist
45from torch.distributed.checkpoint.stateful import Stateful
46from torch.optim.lr_scheduler import LRScheduler
47from torch.utils.data import Dataset
48from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerBase
49from transformers.modeling_outputs import ModelOutput
51from hyper_parallel import DTensor, HSDPModule, SkipDTensorDispatch
52from hyper_parallel.compile.compiler import GraphCompiler
53from hyper_parallel.core.tensor_parallel import loss_parallel
54from hyper_parallel.core.utils import clip_grad_norm_
55from hyper_parallel.trainer.config import (
56 OptimizerSwapConfig,
57 TrainerConfig,
58 normalize_distributed_setup_overrides,
59 save_configs,
60)
61from hyper_parallel.data.dataset_logging import enable_dataset_logging
62from hyper_parallel.data.batching import build_dataloader
63from hyper_parallel.data.parallel import DataLoaderParallelContext, create_dataloader_parallel_context
64from hyper_parallel.trainer.runtime.distributed import (
65 all_reduce,
66 get_global_rank_safe,
67 get_local_rank_safe,
68 get_world_size_safe,
69)
70from hyper_parallel.trainer.runtime.distributed import (
71 create_distributed_setup_from_config,
72 destroy_process_group,
73 initialize_distributed,
74)
75from hyper_parallel.trainer.runtime.logging import setup_logging
76from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
77from hyper_parallel.trainer.runtime.metrics import mean_global_loss
78from hyper_parallel.components.functional.aux_loss import aux_loss_scale_context, bind_aux_loss_scale
79from hyper_parallel.components.modules.moe_utils import MoEMonitorCallback
80from hyper_parallel.models._transformers.loss_parallel import causal_lm_loss_parallel
81from hyper_parallel.models._transformers.model_builder import _validate_optimize_dtype
82from hyper_parallel.models.external_state import get_model_external_state
83from hyper_parallel.components.losses.model_output import ModelOutputLoss
84from hyper_parallel.core.optimizer import (
85 ChainedOptimizer,
86 SwapOptimizerConfig,
87 swap_optimizer,
88)
89from hyper_parallel.components.optim.mixed_precision_optimizer import (
90 Float16OptimizerWithFloat16Params,
91 MixedPrecisionOptimizer,
92)
93from hyper_parallel.trainer.runtime import fsdp as fsdp_runtime
94from hyper_parallel.trainer.runtime import model_integration as model_integration_runtime
95from hyper_parallel.trainer.runtime.data_iterator import HyperIter
96from hyper_parallel.trainer.runtime.logging import enable_third_party_logging
97from hyper_parallel.trainer.runtime.memory import empty_cache, print_device_mem_info
98from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
99from hyper_parallel.trainer.runtime.random import enable_high_precision_for_bf16, set_seed
100from hyper_parallel.trainer.runtime.device import ( # pylint: disable=syntax-error
101 get_device_type,
102 get_torch_device,
103 synchronize,
104)
106from hyper_parallel.trainer.callbacks import (
107 EnvironMeterCallback,
108 EvaluateCallback,
109 GarbageCollectionCallback,
110 LoggingCallback,
111 ProfilerCallback,
112 StressDetectCallback,
113 TqdmCallback,
114 CheckpointerCallback,
115 TrainerState,
116)
118logger = logging.getLogger(__name__)
120if TYPE_CHECKING:
121 from hyper_parallel.data.text.chat_template import ChatTemplate
122 from hyper_parallel.distributed_data import BalancingAlgorithm, CostModel
123 from hyper_parallel.distributed_data.schema import SampleMetadata
126def _to_swap_optimizer_config(config: OptimizerSwapConfig) -> SwapOptimizerConfig:
127 """Translate the trainer-facing swap section into the core swap config.
129 Args:
130 config: Resolved ``optimizer.swap`` configuration section.
132 Returns:
133 The backend-neutral swap configuration consumed by ``swap_optimizer``.
134 """
135 return SwapOptimizerConfig(
136 swap_times=config.swap_times,
137 state_keys=config.state_keys,
138 min_numel=config.min_numel,
139 # ``None`` means "leave it to the backend default", not "off": the core
140 # config declares ``packed_swap: bool = True``, so pass that default
141 # through rather than letting ``None`` read as a hard False.
142 packed_swap=True if config.packed_swap is None else config.packed_swap,
143 )
146def _attach_optimizer_swap(optimizer: Any, config: OptimizerSwapConfig) -> Any:
147 """Attach optimizer-state swap to every Adam/AdamW leaf of ``optimizer``.
149 Every YAML optimizer target returns a ``ChainedOptimizer`` of leaf
150 optimizers -- also for a single AdamW -- and the fp32 main-parameter wrapper
151 keeps the same leaf containers. The swap runtime wraps one concrete
152 Adam/AdamW and delegates everything else, so it is attached to the leaves
153 and the outer object keeps dispatching.
155 Args:
156 optimizer: Optimizer built by the configured target, optionally already
157 wrapped for fp32 main parameters.
158 config: Resolved ``optimizer.swap`` configuration section.
160 Returns:
161 The optimizer carrying swap-wrapped leaves. A leaf the swap runtime does
162 not support fails here, before the first training step.
163 """
164 swap_config = _to_swap_optimizer_config(config)
165 chained = optimizer.optimizer if isinstance(optimizer, MixedPrecisionOptimizer) else optimizer
166 if not isinstance(chained, ChainedOptimizer):
167 return swap_optimizer(optimizer, swap_config)
169 names = list(chained.optimizers_dict)
170 leaves = [swap_optimizer(chained.optimizers_dict[name], swap_config) for name in names]
171 # In place: the fp32 main-parameter wrapper aliases both containers.
172 chained.chained_optimizers[:] = leaves
173 for name, leaf in zip(names, leaves):
174 chained.optimizers_dict[name] = leaf
175 return optimizer
178class BaseTrainer(Stateful, ABC):
179 """
180 Base trainer class for distributed model training.
182 This class provides the core training infrastructure including:
183 - Distributed initialization and parallelism setup
184 - Model, optimizer, and scheduler initialization
185 - Training step execution with gradient accumulation
186 - Checkpointing and fault tolerance
187 - Metrics logging
189 Subclasses can override the following methods to customize behavior:
190 - `_build_distributed_setup()`: Connect AutoModel DistributedSetup/MeshContext
191 - `_build_model()`: Connect AutoModel from_pretrained/parallelize
192 - `_build_model_assets()`: Build tokenizer/processor assets
193 - `_build_data_transform()`: Build modality-specific sample transforms
194 - `forward_backward_step()`: Customize forward/backward logic
195 - `train_step()`: Customize training step execution
196 - `train()`: Train the model
198 Callback Hooks:
199 The trainer calls callback methods at various stages:
200 - evaluate_callback: evaluation callback
201 - trace_callback: tracing callback (meter, wandb, tqdm, profile)
202 - checkpoint_callback: checkpointing callback
203 """
205 # Core configs
206 config: TrainerConfig
207 device: torch.device
208 # Attached at runtime by the model-integration pipeline (the Omni dataset builder sets
209 # ``data_transform``, the integration adapter attaches ``get_batch`` and
210 # ``num_micro_batches``): bare annotations, which assign nothing, so runtime behaviour is
211 # unchanged -- they exist so the type checker sees the contract this class reads.
212 data_transform: Any
213 get_batch: Any
214 num_micro_batches: int
216 # AutoModel distributed setup.
217 # These objects have one owner: ``../build_options.py`` and
218 # ``../components/distributed/mesh.py``. Trainer code consumes them but must
219 # not create a second DeviceMesh or a VeOmni ParallelState.
220 dist_info: Any
221 distributed_setup: Any
222 mesh_context: Any
223 device_mesh: Any
224 moe_mesh: Any
226 # Data
227 train_dataset: Dataset
228 data_transform: Any
229 collate_fn: Any
230 get_batch: Any
231 dataloader_context: DataLoaderParallelContext | None = None
232 train_dataloader: Any = None
233 num_micro_batches: int
235 # Model
236 model: PreTrainedModel = None
237 model_parts: Optional[List[torch.nn.Module]] = None
238 hsdp_model_parts: List[torch.nn.Module] = []
239 fsdp_runtime_diagnostics: Optional[fsdp_runtime.FSDPRuntimeDiagnostics] = None
240 model_integration: Optional[model_integration_runtime.ModelIntegrationSession] = None
241 model_config: PretrainedConfig = PretrainedConfig()
242 tokenizer: PreTrainedTokenizerBase = None
243 processor: Any = None
244 chat_template: Optional[ChatTemplate] = None
245 model_assets: List[Any] = []
247 # Training components
248 optimizer = None
249 lr_scheduler: LRScheduler = None
250 loss_fn: Optional[torch.nn.Module] = None
252 # Training context
253 model_fwd_context: Any
254 model_bwd_context: Any
256 # # Runtime metrics, controlled by trace_callback
257 # environ_meter: helper.EnvironMeter # see in trace_callback.EnvironMeterCallback
258 # step_env_metrics: Dict[str, Any] # mfu, flops, tokens, etc
259 # step_train_metrics: Dict[str, Any] # loss, grad_norm, lr, etc
261 # # Checkpointer
262 # checkpointer: CheckpointerBase # see in checkpoint_callback.CheckpointerCallback
264 # Callback system
265 state: TrainerState
267 # Default seed when training.seed is omitted.
268 default_seed: int = 42
270 def __init__(self, config: TrainerConfig) -> None:
271 """
272 Initialize the trainer.
274 Args:
275 args: Global Arguments
276 Should have attributes: model, data, train
277 model: ModelArguments
278 data: DataArguments
279 train: TrainingArguments
280 """
282 self.config: TrainerConfig = config
284 # General setup owns distributed initialization as its first stage.
285 # ``_build_distributed_setup`` remains a reserved backend hook; the
286 # trainer only defines its contract here.
287 self._setup()
288 memory_profiler.reset(
289 self.config.memory,
290 global_rank=self.global_rank,
291 tp_rank=self.mesh.tp_rank,
292 dp_rank=self.mesh.dp_rank,
293 )
295 # Reserved atomic model-build interface.
296 #
297 # Owner:
298 # - ``../_transformers/auto_model.py::HyperAutoModel*.from_pretrained``
299 # - ``../_transformers/model_builder.py`` for materialization,
300 # checkpoint loading, PEFT/quantization, and parallelize.
301 #
302 # Contract:
303 # - consume the ``DistributedSetup`` created above;
304 # - set ``model`` and the resolved checkpoint ``model_config``;
305 # - return a materialized, weight-loaded, already-parallelized model;
306 # - never call a second trainer-side ``build_parallelize_model``.
307 self._build_model()
308 self._build_loss()
310 # Build trainer-owned data components after the model finalizes parameters and sharding.
311 self._build_model_assets()
312 self._build_data_transform()
313 self._build_dataset()
314 self._build_collate_fn()
315 self._build_dataloader()
316 self._compute_train_iters()
318 self._build_optimizer()
319 self.model_integration.attach_optimizer(self.optimizer)
320 self._build_lr_scheduler()
321 self._build_training_context()
322 self._init_callbacks()
324 def _setup(self):
325 """Initialize logging, distributed state, and the local device."""
326 # log args
327 setup_logging()
329 # Initialize distributed process group
330 backend = getattr(self.config.training, "backend", "nccl")
331 if backend == "hccl" and self.config.training.enable_full_determinism:
332 # HCCL reads this when creating communicators, before set_seed runs below.
333 os.environ["HCCL_DETERMINISTIC"] = "true"
334 initialize_distributed(backend=backend)
336 self.local_rank = get_local_rank_safe()
337 self.global_rank = get_global_rank_safe()
338 self.world_size = get_world_size_safe()
339 if self.config.debug.check_dataset:
340 enable_dataset_logging(self.config.debug.check_dataset)
342 if self.config.training.seed is None:
343 self.config.training.seed = self.default_seed
345 if self.global_rank == 0:
346 logger.info(
347 "Trainer config:\n%s",
348 json.dumps(self.config.to_dict(), indent=2),
349 )
351 # init distributed environment
352 device_str = f"{get_device_type()}:{self.local_rank}"
353 get_torch_device().set_device(device_str)
354 self.device = torch.device(device_str)
356 logger.info("Process rank: %s, world size: %s", self.global_rank, self.world_size)
358 # Register distributed_setup; the Trainer-side desugar fills its
359 # normalized plan_overrides/module_replacements (05 §15.2.6) so the
360 # AutoModels build pipeline never sees raw YAML PlanOverride DTOs.
361 self.distributed_setup = create_distributed_setup_from_config(self.config)
362 normalize_distributed_setup_overrides(self.distributed_setup, self.config)
363 self.mesh = self.distributed_setup.mesh_context
364 self.device_mesh = self.mesh.device_mesh
365 self.dp_cp_mesh = self.mesh.dp_cp_mesh
367 # Set random seed
368 set_seed(self.config.training.seed, self.config.training.enable_full_determinism)
370 # Enable high precision for bf16
371 enable_high_precision_for_bf16()
373 # Enable third party logging
374 if self.local_rank == 0:
375 enable_third_party_logging()
377 # Save arguments
378 if self.global_rank == 0:
379 save_configs(self.config, self.config.checkpoint.checkpoint_dir)
381 def _build_model(self) -> None:
382 """Build the model and derive Trainer-owned runtime state."""
383 self.peft_config = self.config.peft
384 self.model = self.config.model.build(
385 distributed_setup=self.distributed_setup,
386 peft_config=self.peft_config,
387 activation_checkpoint=self.config.activation_checkpoint.mode,
388 activation_checkpoint_selection=self.config.activation_checkpoint.selection,
389 swap_inputs=getattr(self.config.activation_checkpoint, "swap_inputs", False),
390 allow_cache_entry_mutation=getattr(
391 self.config.activation_checkpoint, "allow_cache_entry_mutation", False),
392 activation_swap=self.config.activation_swap,
393 compile_config=self.config.compile,
394 # The final dtype is applied inside the atomic build (05 stage-5
395 # item 5); the Trainer no longer patches it afterwards.
396 model_init_dtype=self.config.model_init_dtype,
397 )
398 _validate_optimize_dtype(
399 self.model,
400 fp32_main_params=self.config.optimizer.fp32_main_params,
401 )
402 self.model_config = self.model.config
403 if self.global_rank == 0:
404 logger.info(
405 "Model config:\n%s",
406 json.dumps(self.model.config.to_dict(), indent=2),
407 )
408 if self.config.accelerator.loss_parallel:
409 self.model.loss_function = causal_lm_loss_parallel
410 model_parts = getattr(self.model, "parts", None)
411 self.model_parts = list(model_parts) if model_parts is not None else [self.model]
412 self.hsdp_model_parts = [
413 model_part
414 for model_part in self.model_parts
415 if isinstance(model_part, HSDPModule)
416 ]
417 self.model_integration = model_integration_runtime.build_model_integration_session(
418 self.model,
419 self.config,
420 self.global_rank,
421 self.mesh,
422 )
423 self.fsdp_runtime_diagnostics = (
424 fsdp_runtime.build_fsdp_runtime_diagnostics(self.hsdp_model_parts)
425 if getattr(self.config.debug, "check_fsdp_runtime", False)
426 else None
427 )
429 def _build_loss(self) -> None:
430 """Build the configured loss module or use the model-output default."""
431 loss_fn = (
432 self.config.loss_fn.build()
433 if self.config.loss_fn is not None
434 else ModelOutputLoss()
435 )
436 if not isinstance(loss_fn, torch.nn.Module):
437 raise ValueError("config.loss_fn must build a torch.nn.Module")
438 self.loss_fn = loss_fn
439 # Model-integrated objectives may need to intercept the model before its
440 # ordinary forward materializes an otherwise avoidable output tensor.
441 loss_model_binding_hook = getattr(self.loss_fn, "bind_model", None)
442 if callable(loss_model_binding_hook):
443 loss_model_binding_hook(self.model, self.distributed_setup)
445 def _build_model_assets(self) -> None:
446 """Require a concrete Trainer to build modality-specific model assets."""
447 raise NotImplementedError("Concrete Trainer must implement _build_model_assets")
449 def _build_data_transform(self) -> None:
450 """Require a concrete Trainer to build its sample transform."""
451 raise NotImplementedError("Concrete Trainer must implement _build_data_transform")
453 def _build_dataset(self) -> None:
454 """Build and assign one Dataset or a train-validation-test tuple."""
455 data_config = getattr(self.config.dataset, "data_config", {})
456 self.dataloader_context = create_dataloader_parallel_context(
457 self.mesh,
458 data_index_cache=bool(data_config.get("data_index_cache", False)),
459 shared_storage=not bool(data_config.get("no_shared_storage", False)),
460 tp_source_rank=None if self.config.dataloader.load_balance is not None else 0,
461 )
462 dataset_result = self.config.dataset.build(
463 transform=self.data_transform,
464 tokenizer=self.tokenizer,
465 mesh_context=self.mesh,
466 training_config=self.config.training,
467 sampler_type=self.config.dataloader.sampler_type,
468 dataloader_context=self.dataloader_context,
469 )
471 if isinstance(dataset_result, tuple) and len(dataset_result) == 3:
472 self.train_dataset, self.valid_dataset, self.test_dataset = dataset_result
473 return
475 self.train_dataset = dataset_result
476 self.valid_dataset = self.test_dataset = None
478 def _build_collate_fn(self) -> None:
479 """Require a concrete Trainer to build its micro-batch collator."""
480 raise NotImplementedError("Concrete Trainer must implement _build_collate_fn")
482 def _build_dataloader(
483 self,
484 *,
485 metadata_fn: Callable[[Any], SampleMetadata] | None = None,
486 cost_model: CostModel | None = None,
487 balancing_algorithm: BalancingAlgorithm | None = None,
488 ) -> None:
489 """Build and assign train, validation, and test dataloaders."""
490 split_names = ("train", "valid", "test")
491 datasets = tuple(getattr(self, f"{split_name}_dataset", None) for split_name in split_names)
492 dataloaders, batch_samplers = build_dataloader(
493 self.config.dataloader,
494 datasets=datasets,
495 collate_fn=self.collate_fn,
496 mesh_context=self.mesh,
497 training_config=self.config.training,
498 data_config=getattr(self.config.dataset, "data_config", {}),
499 max_seq_len=getattr(self.data_transform, "max_seq_len", None),
500 default_seed=self.default_seed,
501 metadata_fn=metadata_fn,
502 model_config=self.model_config,
503 cost_model=cost_model,
504 balancing_algorithm=balancing_algorithm,
505 )
506 for split_name, dataloader, batch_sampler in zip(split_names, dataloaders, batch_samplers):
507 setattr(self, f"{split_name}_dataloader", dataloader)
508 setattr(self, f"{split_name}_batch_sampler", batch_sampler)
510 def attach_model_integration_data_pipeline(self) -> None:
511 """Attach the final data pipeline through one shared Trainer path."""
512 self.model_integration.attach_data_pipeline(
513 runtime_adapter=getattr(self.get_batch, "runtime_input_adapter", None),
514 get_batch=self.get_batch,
515 model_assets=tuple(self.model_assets),
516 )
518 def _compute_train_iters(self) -> None:
519 """Resolve the run limit and the optimizer steps available per Dataset epoch."""
520 training_config = self.config.training
521 if training_config.train_iters is not None and training_config.train_samples is not None:
522 raise ValueError("configure only one of training.train_iters and training.train_samples")
524 if training_config.train_iters is None and training_config.train_samples is None:
525 raise ValueError("configure one of training.train_iters or training.train_samples")
527 train_iters = training_config.train_iters
528 if training_config.train_samples is not None:
529 train_iters = training_config.train_samples // training_config.global_batch_size
531 if train_iters <= 0:
532 raise ValueError("training configuration must produce at least one train iteration")
534 try:
535 micro_batches_per_epoch = len(self.train_dataloader)
536 except TypeError:
537 micro_batches_per_epoch = 0
539 if micro_batches_per_epoch:
540 train_steps = micro_batches_per_epoch // self.num_micro_batches
541 if train_steps <= 0:
542 raise ValueError("train DataLoader must provide one complete optimizer step per epoch")
543 else:
544 # Streaming and dynamic DataLoaders define an epoch only when their
545 # source iterator is exhausted.
546 train_steps = train_iters
548 self.train_iters = train_iters
549 self.train_steps = train_steps
550 self.train_epochs = (self.train_iters + self.train_steps - 1) // self.train_steps
552 def _build_optimizer(self) -> None:
553 """Build the configured optimizer with the runtime model context."""
554 config: TrainerConfig = self.config
555 external = get_model_external_state(self.model)
556 if external is not None:
557 def wrap_dense(dense_raw: Any) -> Any:
558 """Wrap dense parameters in FP32 main and optional state swap.
560 Args:
561 dense_raw: Constructed dense optimizer.
562 """
563 dense = (
564 Float16OptimizerWithFloat16Params(dense_raw, self.model)
565 if config.optimizer.fp32_main_params else dense_raw
566 )
567 if config.optimizer.swap.enabled:
568 dense = _attach_optimizer_swap(dense, config.optimizer.swap)
569 return dense
571 self.optimizer = external.build_optimizer(config.optimizer, wrap_dense=wrap_dense)
572 return
573 optimizer = config.optimizer.target.build(model=self.model).get_optimizer()
574 self.optimizer = (
575 Float16OptimizerWithFloat16Params(optimizer, self.model)
576 if config.optimizer.fp32_main_params
577 else optimizer
578 )
579 # Swap is attached after the fp32 main-parameter wrap so leaf optimizers
580 # already reference their final parameters, and the lr scheduler built
581 # right after reads the wrapped ``self.optimizer.param_groups``.
582 if config.optimizer.swap.enabled:
583 self.optimizer = _attach_optimizer_swap(self.optimizer, config.optimizer.swap)
585 def _build_lr_scheduler(self):
586 config: TrainerConfig = self.config
587 scheduler_iters = config.training.lr_scheduler_iters or self.train_iters
588 if scheduler_iters <= 0:
589 raise ValueError("training.lr_scheduler_iters must be positive when configured")
590 lr_scheduler = config.lr_scheduler.build(
591 optimizer=self.optimizer,
592 train_iters=scheduler_iters,
593 )
594 self.lr_scheduler = lr_scheduler.get_lr_scheduler()
596 def _build_training_context(self):
597 """Build training context for distributed training."""
598 if self.config.accelerator.loss_parallel:
599 tp_mesh = self.device_mesh["tp"]
600 self.model_fwd_context = partial(loss_parallel, mesh=tp_mesh)
601 else:
602 self.model_fwd_context = nullcontext()
603 self.model_bwd_context = nullcontext()
604 self.graph_compiler: Optional[GraphCompiler] = None
606 def _build_graph_compiler(self) -> None:
607 """Build the optional joint-graph compiler for this trainer runtime."""
608 if not self.config.compile.selects_graph_compiler():
609 return
610 self.graph_compiler = GraphCompiler(
611 model=self.model,
612 train_fn=self._graph_train_fn,
613 trainer_config=self.config,
614 device=self.device,
615 mesh_context=self.mesh,
616 source_shard_info=getattr(self.model, "_source_shard_info", None),
617 )
619 def _graph_train_fn(
620 self,
621 model: torch.nn.Module,
622 model_inputs: Dict[str, Any],
623 labels: Optional[torch.Tensor],
624 current_token_counts: Dict[str, torch.Tensor],
625 ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
626 """Capture forward/loss with explicit runtime token counts."""
627 model_fwd_context = (
628 self.model_fwd_context()
629 if callable(self.model_fwd_context)
630 else self.model_fwd_context
631 )
632 with model_fwd_context:
633 outputs = model(**model_inputs, use_cache=False)
634 # Counts must be graph inputs: packed lengths and ignored-label counts
635 # can vary even when tensor shapes happen to match the first batch.
636 loss, loss_dict = self.postforward(
637 outputs, labels,
638 current_token_counts=current_token_counts,
639 )
640 return loss, loss_dict
642 def _init_callbacks(self):
643 """Initialize callbacks."""
644 dp_cp_mesh = self.mesh.dp_cp_mesh
645 self.moe_monitor = MoEMonitorCallback(
646 self.model,
647 dp_group=dp_cp_mesh.get_group() if dp_cp_mesh is not None else None,
648 tp_group=self.device_mesh["tp"].get_group() if self.mesh.sequence_parallel else None,
649 )
650 self.state = TrainerState()
651 # Raw AUX is not normalized with the backward loss; average its observations.
652 self.environ_meter_callback = EnvironMeterCallback(self, microbatch_mean_metrics=("router_aux_loss",))
653 self.tqdm_callback = TqdmCallback(self)
654 self.logging_callback = LoggingCallback(self)
655 self.evaluate_callback = EvaluateCallback(self)
656 self.garbage_collection_callback = GarbageCollectionCallback(self)
657 self.profiler_callback = ProfilerCallback(self)
658 self._callbacks = [
659 self.environ_meter_callback,
660 self.logging_callback,
661 self.tqdm_callback,
662 self.evaluate_callback,
663 self.garbage_collection_callback,
664 self.profiler_callback,
665 ]
666 if self.config.stress_detect.op_stress_detect.detection_interval > 0:
667 self.stress_detect_callback = StressDetectCallback(self)
668 self._callbacks.append(self.stress_detect_callback)
669 # Registered to save, to restore, or both --- the two are independent, so
670 # a run that only loads an existing checkpoint still needs the callback.
671 if (
672 self.config.checkpoint.save_ckpt
673 or self.config.checkpoint.save_hf_weights
674 or self.config.checkpoint.restore_from is not None
675 ):
676 self._callbacks.append(CheckpointerCallback(self))
677 else:
678 logger.info(
679 "Checkpointing is inactive (save_ckpt=false, save_hf_weights=false, "
680 "restore_from=None); "
681 "no checkpoint callback registered."
682 )
684 def on_train_begin(self) -> None:
685 """Run all registered callbacks at the start of training."""
686 for callback in self._callbacks:
687 callback.on_train_begin(self.state)
689 def on_train_end(self) -> None:
690 """Run all registered callbacks at the end of training."""
691 for callback in self._callbacks:
692 callback.on_train_end(self.state)
694 def on_epoch_begin(self) -> None:
695 """Run all registered callbacks at the start of an epoch."""
696 for callback in self._callbacks:
697 callback.on_epoch_begin(self.state)
699 def on_epoch_end(self) -> None:
700 """Run all registered callbacks at the end of an epoch."""
701 for callback in self._callbacks:
702 callback.on_epoch_end(self.state)
704 def on_step_begin(self, **kwargs: Any) -> None:
705 """Run all registered callbacks at the start of a training step."""
706 for callback in self._callbacks:
707 callback.on_step_begin(self.state, **kwargs)
709 def on_micro_step_begin(self, micro_batch: Dict[str, Any], **kwargs: Any) -> None:
710 """Run all registered callbacks before one forward-backward micro step.
712 Args:
713 micro_batch: Prepared inputs for the current micro step.
714 **kwargs: Additional callback context.
715 """
716 for callback in self._callbacks:
717 callback.on_micro_step_begin(self.state, micro_batch, **kwargs)
719 def on_micro_step_end(self, loss_dict: Dict[str, torch.Tensor], **kwargs: Any) -> None:
720 """Run all registered callbacks after one forward-backward micro step.
722 Args:
723 loss_dict: Named losses and metrics returned by the completed micro step.
724 **kwargs: Additional callback context.
725 """
726 for callback in self._callbacks:
727 callback.on_micro_step_end(self.state, loss_dict=loss_dict, **kwargs)
729 def on_step_end(
730 self,
731 loss: Optional[float] = None,
732 loss_dict: Optional[Dict[str, Any]] = None,
733 grad_norm: Optional[float] = None,
734 ) -> None:
735 """Run all registered callbacks at the end of a training step.
737 Args:
738 loss: Aggregated loss for the optimizer step.
739 loss_dict: Named loss values for the optimizer step.
740 grad_norm: Gradient norm measured before the optimizer update.
741 """
742 for callback in self._callbacks:
743 callback.on_step_end(self.state, loss=loss, loss_dict=loss_dict, grad_norm=grad_norm)
745 def preforward(self, micro_batch: List[Dict[str, Any]]) -> Dict[str, Any]:
746 """Preprocess micro batches before forward pass.
748 Tensors are moved to ``self.device`` non-blockingly. Nested dicts
749 (e.g. ``multimodal_metadata`` emitted by ``PackingCollator``) are
750 recursed so inner tensor values land on the device too; Python ints
751 / lists / etc. pass through unchanged.
753 Args:
754 micro_batch: Batch fields to move to the training device.
756 Returns:
757 Batch fields with tensors moved to the training device.
758 """
760 def _to_device(v: Any) -> Any:
761 if isinstance(v, torch.Tensor):
762 return v.to(self.device, non_blocking=True)
763 if isinstance(v, dict):
764 return {k: _to_device(vv) for k, vv in v.items()}
765 return v
767 # chunk_mbs_config = getattr(self.config.training, "chunk_mbs_config", None)
768 # self._chunk_mbs_ranges = build_chunk_mbs_ranges(micro_batch, chunk_mbs_config)
769 micro_batch = {k: _to_device(v) for k, v in micro_batch.items()} # type: ignore
770 # if getattr(self, "LOG_SAMPLE", True):
771 # self.LOG_SAMPLE = False
772 return micro_batch
774 def postforward(
775 self, outputs: ModelOutput, labels: Optional[torch.Tensor],
776 *,
777 current_token_counts: Dict[str, torch.Tensor],
778 ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
779 """Postprocess model outputs after forward pass.
781 Args:
782 outputs: Model outputs consumed by the configured loss.
783 labels: Labels associated with the current micro-batch.
784 current_token_counts: Per-loss token counts for the current micro-batch.
786 Returns:
787 Backward loss and named globally aggregated loss values.
788 """
789 local_loss = self.loss_fn(model_output=outputs, labels=labels)
790 # Names are static graph metadata; diagnostics added below are not losses.
791 self.loss_names = tuple(local_loss) if isinstance(local_loss, dict) else ("foundation_loss",)
792 # Every foundation-family objective receives the same external token weight.
793 # Bind each branch before backward can reach its independently injected losses.
794 foundation_losses = (
795 (value for key, value in local_loss.items() if key.split("_loss", maxsplit=1)[0] == "foundation")
796 if isinstance(local_loss, dict) else (local_loss,)
797 )
798 for foundation_loss in foundation_losses:
799 if foundation_loss is not None:
800 bind_aux_loss_scale(foundation_loss)
801 loss_dict: Dict[str, torch.Tensor] = mean_global_loss(
802 local_loss,
803 current_token_counts,
804 device_mesh=self.mesh,
805 )
806 loss = torch.stack(list(loss_dict.values())).sum()
807 # Model-declared extra metrics join the logged mapping only after
808 # the backward loss is fixed: log-only, never part of the loss.
809 extra_metrics = getattr(outputs, "extra_metrics", None)
810 if extra_metrics is None:
811 model_ref = getattr(self, "model", None)
812 extra_metrics = getattr(model_ref, "pending_extra_metrics", None)
813 if extra_metrics is not None:
814 model_ref.pending_extra_metrics = None
815 if extra_metrics:
816 for name, value in extra_metrics.items():
817 loss_dict[name] = value.detach().float()
818 # Log-only outputs share the eager/graph path without entering backward loss.
819 moe_monitor = getattr(self, "moe_monitor", None)
820 if moe_monitor is not None:
821 aux_loss = moe_monitor.last_mean_aux_loss
822 if aux_loss is not None:
823 loss_dict["router_aux_loss"] = aux_loss
824 return loss, loss_dict
826 def forward_backward_step(
827 self,
828 micro_batch: dict[str, torch.Tensor],
829 loss_inputs: Optional[dict[str, torch.Tensor]] = None,
830 ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
831 """Run forward and backward computation for one micro batch.
833 Args:
834 micro_batch: Inputs forwarded to the model.
835 loss_inputs: Loss-only labels and masks from the get-batch runtime. If
836 omitted, the legacy single-dictionary path is preserved.
838 Returns:
839 Backward loss and the named globally aggregated loss mapping.
840 """
841 for name, count in self.current_token_counts.items():
842 self.step_token_counts[name] = self.step_token_counts.get(name, 0) + count.detach()
843 channel_loss_callback = getattr(self, "channel_loss_callback", None)
844 micro_step_context = (
845 channel_loss_callback.micro_step_context(self.state, micro_batch)
846 if channel_loss_callback is not None
847 else nullcontext()
848 )
849 with micro_step_context, aux_loss_scale_context():
850 micro_batch = self.preforward(micro_batch)
851 # TextTrainer passes loss-only fields separately. Keep the legacy
852 # BaseTrainer/VLM one-dictionary call contract working as well.
853 if loss_inputs is None:
854 loss_inputs = {
855 name: micro_batch[name]
856 for name in ("labels", "shift_labels", "loss_mask")
857 if name in micro_batch
858 }
859 else:
860 loss_inputs = self.preforward(loss_inputs)
861 labels = loss_inputs.get("labels", micro_batch.get("labels"))
863 # A model-integrated loss can translate the public batch fields into
864 # model-family protocol arguments before the model forward starts.
865 prepare_model_inputs = getattr(
866 self.loss_fn,
867 "prepare_model_inputs",
868 None,
869 )
870 if callable(prepare_model_inputs):
871 micro_batch = prepare_model_inputs(micro_batch, loss_inputs)
872 if channel_loss_callback is not None:
873 channel_loss_callback.strip_model_inputs(micro_batch)
874 self.model_integration.record_batch(micro_batch, loss_inputs)
876 if self.graph_compiler is not None:
877 runtime_adapter = getattr(getattr(self, "get_batch", None), "runtime_input_adapter", None)
878 if runtime_adapter is not None:
879 micro_batch = runtime_adapter.prepare_graph_inputs(micro_batch)
880 if self.config.training.empty_cache_before_backward:
881 logger.warning(
882 "training.empty_cache_before_backward is ignored in graph compiler mode "
883 "because forward and backward execute in one joint graph"
884 )
885 # The traced ``loss_dict`` values flow back as extra graph
886 # outputs (emitted after the gradients), so the logging
887 # metrics keep per-key parity with the eager path.
888 loss, loss_dict = self.graph_compiler.forward_backward(
889 model_inputs=micro_batch,
890 labels=labels,
891 current_token_counts=self.current_token_counts,
892 )
893 else:
894 model_fwd_context = (
895 self.model_fwd_context()
896 if callable(self.model_fwd_context)
897 else self.model_fwd_context
898 )
899 with model_fwd_context:
900 outputs: ModelOutput = self.model(**micro_batch, use_cache=False)
902 # with use_parallel_state("base"):
903 loss, loss_dict = self.postforward(
904 outputs, labels,
905 current_token_counts=self.current_token_counts,
906 )
907 prefetch_plan = getattr(self.train_dataloader, "prefetch_plan", None)
908 if callable(prefetch_plan):
909 prefetch_plan() # pylint: disable=not-callable # Guarded dynamic hook.
910 # The loss graph owns everything required for backward. Releasing
911 # the model output here avoids retaining large vocabulary logits
912 # until the whole backward pass finishes.
913 del outputs
914 if self.config.training.empty_cache_before_backward:
915 empty_cache()
917 # Backward pass
918 with self.model_bwd_context:
919 loss.backward()
920 prefetch = getattr(self.train_dataloader, "prefetch", None)
921 if callable(prefetch):
922 prefetch() # pylint: disable=not-callable # Guarded dynamic hook.
924 del micro_batch
925 return loss, loss_dict
927 def model_reshard(self, micro_step: int, num_micro_steps: int) -> None:
928 """Reshard model after backward pass; policy lives in ``runtime/fsdp.py``.
930 Args:
931 micro_step: Zero-based micro-step index.
932 num_micro_steps: Number of micro-steps in the optimizer step.
933 """
934 fsdp_runtime.model_reshard(self.hsdp_model_parts, self.config.fsdp_config, micro_step, num_micro_steps)
936 def _configure_fsdp_gradient_sync(self, micro_step: int, num_micro_steps: int):
937 """Configure FSDP gradient synchronization; policy lives in ``runtime/fsdp.py``."""
938 fsdp_runtime.configure_fsdp_gradient_sync(
939 self.hsdp_model_parts,
940 self.config.fsdp_config,
941 self.mesh.dp_replicate_size,
942 micro_step,
943 num_micro_steps,
944 )
946 def configure_fsdp_gradient_sync(self, micro_step: int, num_micro_steps: int) -> None:
947 """Configure FSDP gradient synchronization for an external training loop.
949 Args:
950 micro_step: Zero-based micro-step index.
951 num_micro_steps: Number of micro-steps in the optimizer step.
952 """
953 self._configure_fsdp_gradient_sync(micro_step, num_micro_steps)
955 def begin_fsdp_runtime_diagnostics(self, micro_step: int) -> None:
956 """Begin observing one FSDP forward/backward micro-step when enabled.
958 Args:
959 micro_step: Microbatch index.
960 """
961 if self.fsdp_runtime_diagnostics is not None:
962 self.fsdp_runtime_diagnostics.begin_micro_step(
963 self.state.global_step,
964 micro_step,
965 )
967 def validate_fsdp_runtime_before_optimizer(self) -> None:
968 """Fail before optimizer updates when FSDP prefetch state is unsafe."""
969 if self.fsdp_runtime_diagnostics is not None:
970 self.model_integration.record_fsdp_trace(
971 self.fsdp_runtime_diagnostics.trace_evidence()
972 )
973 self.fsdp_runtime_diagnostics.validate_before_optimizer_step()
975 def begin_loss_step(self) -> None:
976 """Reset scalar counters before reading the next optimizer step."""
977 self.step_token_counts = {}
979 def _loss_denominator(self) -> float:
980 """Return the optimizer-step valid-token count across DP/CP."""
981 family = self.loss_names[0].split("_loss", maxsplit=1)[0]
982 count = self.step_token_counts[f"{family}_tokens"].detach().clone()
983 mesh = self.mesh.dp_cp_mesh
984 if mesh is not None and mesh.size() > 1:
985 count = all_reduce(count.to(self.device), op="sum", group=mesh.get_group())
986 else:
987 count = count.item()
988 return max(count, 1)
990 @torch.no_grad()
991 def loss_normalize(self, loss: float, loss_dict: Dict[str, float], *, denominator: float | None = None) -> float:
992 """Normalize Trainer-owned gradients and loss metrics before clipping.
994 External providers expose their owned parameters through ``parameters``
995 and normalize those gradients in ``prepare_optimizer_step`` instead.
996 """
997 if denominator is None:
998 denominator = self._loss_denominator()
999 external = get_model_external_state(self.model)
1000 excluded = {id(param) for param in external.parameters.values()} if external is not None else set()
1001 with SkipDTensorDispatch():
1002 for parameter in self.model.parameters():
1003 if id(parameter) in excluded:
1004 continue
1005 grad = getattr(parameter, "main_grad", None)
1006 if grad is None:
1007 grad = parameter.grad
1008 if grad is not None:
1009 if isinstance(grad, DTensor):
1010 grad = grad.to_local()
1011 grad.div_(denominator)
1012 for name in self.loss_names:
1013 loss_dict[name] /= denominator
1014 return loss / denominator
1016 def prepare_optimizer_step(self, loss: float, loss_dict: Dict[str, float]) -> tuple[float, Any]:
1017 """Normalize the step objective before clipping and optimizer preparation."""
1018 denominator = self._loss_denominator()
1019 loss = self.loss_normalize(loss, loss_dict, denominator=denominator)
1020 return loss, self.clip_grad_norm(gradient_divisor=denominator)
1022 def clip_grad_norm(self, *, gradient_divisor: float = 1.0) -> Any:
1023 """Run pre-optimizer diagnostics and clipping in their required order."""
1024 self.model_integration.after_backward_before_clip()
1025 self.validate_fsdp_runtime_before_optimizer()
1026 max_grad_norm = self.config.training.max_grad_norm
1027 external = get_model_external_state(self.model)
1028 if external is not None:
1029 preparation = external.prepare_optimizer_step(max_norm=max_grad_norm, gradient_divisor=gradient_divisor)
1030 self.model_integration.after_clip(preparation.post_clip_norm)
1031 return preparation.global_norm
1032 grad_norm: Any = 0.0
1033 post_clip_norm = None
1034 compiler = getattr(self, "graph_compiler", None)
1035 clip_kwargs = {"parameter_layouts": compiler.parameter_layouts} if compiler is not None else {}
1036 if max_grad_norm > 0:
1037 grad_norm = clip_grad_norm_(self.model, max_grad_norm, **clip_kwargs)
1038 if self.model_integration.runtime_enabled:
1039 grad_norm_value = (
1040 float(grad_norm.item())
1041 if isinstance(grad_norm, torch.Tensor)
1042 else float(grad_norm)
1043 )
1044 clip_coefficient = min(
1045 float(max_grad_norm) / (grad_norm_value + 1.0e-6),
1046 1.0,
1047 )
1048 post_clip_norm = grad_norm_value * clip_coefficient
1049 elif self.model_integration.runtime_enabled:
1050 grad_norm = clip_grad_norm_(self.model, float("inf"), **clip_kwargs)
1051 post_clip_norm = grad_norm
1052 self.model_integration.after_clip(post_clip_norm)
1053 return grad_norm
1055 def step_optimizers_and_schedulers(self) -> None:
1056 """Step every optimizer and scheduler, then expose the final state."""
1057 optimizers = self.optimizer if isinstance(self.optimizer, list) else [self.optimizer]
1058 for optimizer in optimizers:
1059 # Keep zeros_like dispatched so lazily created optimizer states retain DTensor metadata.
1060 with SkipDTensorDispatch(no_skip={torch.zeros_like}):
1061 optimizer.step()
1062 optimizer.zero_grad()
1063 moe_monitor = getattr(self, "moe_monitor", None)
1064 if moe_monitor is not None:
1065 moe_monitor.on_step_end()
1066 self.model_integration.after_optimizer()
1068 schedulers = (
1069 self.lr_scheduler
1070 if isinstance(self.lr_scheduler, list)
1071 else ([self.lr_scheduler] if self.lr_scheduler is not None else [])
1072 )
1073 for scheduler in schedulers:
1074 scheduler.step()
1076 def _end_model_integration_step(self, metrics: Dict[str, Any]) -> None:
1077 """Persist structured step metrics and global input identity."""
1078 if not self.model_integration.runtime_enabled:
1079 return
1080 resolved_metrics = dict(metrics)
1081 optimizers = self.optimizer if isinstance(self.optimizer, list) else [self.optimizer]
1082 learning_rates = []
1083 for optimizer in optimizers:
1084 inner = getattr(optimizer, "optimizer", optimizer)
1085 leaves = getattr(inner, "chained_optimizers", (inner,))
1086 for leaf in leaves:
1087 learning_rates.extend(
1088 group.get("lr")
1089 for group in getattr(leaf, "param_groups", ())
1090 if group.get("lr") is not None
1091 )
1092 if learning_rates:
1093 resolved_metrics["lr"] = float(learning_rates[0])
1094 token_counts = getattr(self, "step_token_counts", {})
1095 token_values = []
1096 for value in token_counts.values():
1097 token_values.append(float(value.item() if torch.is_tensor(value) else value))
1098 if token_values:
1099 resolved_metrics["tokens"] = sum(token_values)
1100 resolved_metrics["samples"] = self.config.training.global_batch_size
1101 self.model_integration.end_step(resolved_metrics)
1103 def train_step(
1104 self,
1105 data_iterator: Any,
1106 ) -> Dict[str, float]:
1107 """Execute one optimizer update from the next dataloader batch.
1109 Args:
1110 data_iterator: Iterator providing the next micro-batches.
1112 Returns:
1113 Aggregated loss and gradient norm.
1114 """
1115 memory_profiler.step()
1116 micro_batches: List[Dict[str, Any]] = next(data_iterator)
1117 self.state.global_step += 1
1119 self.on_step_begin(micro_batches=micro_batches)
1120 self.model_integration.begin_step(self.state.global_step)
1122 # Forward and backward for each micro batch
1123 synchronize()
1125 total_loss = 0.0
1126 total_loss_dict = defaultdict(int)
1128 # token num for fixed_ce_loss in postforward
1129 self.begin_loss_step()
1130 num_micro_steps = len(micro_batches)
1131 # forward and backward pass with gradient_accumulationsteps
1132 for micro_step, micro_batch in enumerate(micro_batches):
1133 self.model_reshard(micro_step, num_micro_steps)
1134 self._configure_fsdp_gradient_sync(micro_step, num_micro_steps)
1135 self.begin_fsdp_runtime_diagnostics(micro_step)
1136 loss: torch.Tensor
1137 loss_dict: Dict[str, torch.Tensor]
1138 # token num for fixed_ce_loss in postforward
1139 self.current_token_counts = count_loss_token(micro_batch)
1140 loss, loss_dict = self.forward_backward_step(micro_batch)
1141 self.on_micro_step_end(loss_dict)
1143 total_loss += loss.item()
1144 for k, v in loss_dict.items():
1145 total_loss_dict[k] += v.item()
1147 total_loss, grad_norm = self.prepare_optimizer_step(total_loss, total_loss_dict)
1148 self.step_optimizers_and_schedulers()
1150 grad_norm_value = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else float(grad_norm)
1151 self._end_model_integration_step({"loss": total_loss, "grad_norm": grad_norm_value})
1152 self.on_step_end(loss=total_loss, loss_dict=total_loss_dict, grad_norm=grad_norm_value)
1153 return {
1154 "loss": total_loss,
1155 "grad_norm": grad_norm_value,
1156 }
1158 def abort_profilers(self) -> None:
1159 """Release process-local profiling resources while preserving a training error."""
1160 try:
1161 memory_profiler.abort()
1162 # Cleanup must also preserve the training error if interrupted.
1163 except BaseException: # pylint: disable=broad-exception-caught
1164 logger.exception("Failed to stop memory profiling while aborting training.")
1165 profiler_callback = getattr(self, "profiler_callback", None)
1166 if profiler_callback is not None:
1167 try:
1168 profiler_callback.close()
1169 # Cleanup must also preserve the training error if interrupted.
1170 except BaseException: # pylint: disable=broad-exception-caught
1171 logger.exception("Failed to stop performance profiling while aborting training.")
1173 def destroy_distributed(self) -> None:
1174 """Synchronize all ranks and tear down the distributed process group."""
1175 if not dist.is_available() or not dist.is_initialized():
1176 return
1178 empty_cache()
1179 dist.barrier()
1181 synchronize()
1182 destroy_process_group()
1184 def train(self) -> None:
1185 """Run the configured training loop."""
1186 try:
1187 config: TrainerConfig = self.config
1188 self.on_train_begin()
1189 self.data_iterator = HyperIter(
1190 self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
1191 )
1192 logger.info(
1193 "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
1194 self.local_rank, self.state.global_step, self.train_iters, self.state.epoch, self.train_epochs,
1195 )
1197 try:
1198 start_epoch = self.state.epoch
1199 for epoch in range(start_epoch, self.train_epochs):
1200 if epoch != start_epoch:
1201 self.train_dataloader.set_epoch(epoch)
1202 self.data_iterator = HyperIter(
1203 self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
1204 )
1205 self.state.epoch = epoch
1207 self.on_epoch_begin()
1209 start_step = self.state.global_step - epoch * self.train_steps
1210 train_steps = min(self.train_steps, self.train_iters - epoch * self.train_steps)
1211 for _ in range(start_step, train_steps):
1212 try:
1213 self.train_step(self.data_iterator)
1214 except StopIteration:
1215 logger.info(
1216 "epoch:%s Dataloader finished with drop_last %s",
1217 epoch,
1218 config.dataloader.drop_last,
1219 )
1220 break
1222 self.on_epoch_end()
1223 self.state.epoch = epoch + 1
1225 print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")
1227 if config.dataloader.use_background_prefetcher:
1228 self.data_iterator.stop()
1230 memory_profiler.stop()
1231 self.on_train_end()
1232 finally:
1233 if config.dataloader.use_background_prefetcher:
1234 self.data_iterator.stop()
1236 except BaseException:
1237 try:
1238 memory_profiler.stop()
1239 # Snapshot export must preserve the training error even if interrupted.
1240 except BaseException: # pylint: disable=broad-exception-caught
1241 logger.exception("Failed to save memory snapshot after training error.")
1242 finally:
1243 self.abort_profilers()
1244 raise
1246 synchronize()
1248 self.destroy_distributed()