Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / text_trainer.py: 39%
172 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.
15"""Text Trainer assembled from the shared BaseTrainer stages."""
17__all__ = ["TextTrainer"]
19from collections import defaultdict
20from typing import Any, Dict
22import torch # pylint: disable=forbidden-backend-import
24from hyper_parallel.data.batching import calculate_num_micro_batches
25from hyper_parallel.data.text import build_chat_template
26from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
27from hyper_parallel.trainer.runtime.logging import create_logger
28from hyper_parallel.trainer.runtime.memory import print_device_mem_info
29from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
30from hyper_parallel.trainer.runtime.device import synchronize
31from hyper_parallel.trainer.base import BaseTrainer
32from hyper_parallel.trainer.config import TrainerConfig
34logger = create_logger(__name__)
37class TextTrainer:
38 """Compose the text training runtime from explicit BaseTrainer stages."""
40 base: BaseTrainer
42 def __init__(self, config: TrainerConfig) -> None:
43 """Build the text Trainer in the same explicit order as VeOmni."""
44 self.base = BaseTrainer.__new__(BaseTrainer)
45 self.base.config = config
47 self.base._setup()
48 memory_profiler.reset(
49 self.base.config.memory,
50 global_rank=self.base.global_rank,
51 tp_rank=self.base.mesh.tp_rank,
52 dp_rank=self.base.mesh.dp_rank,
53 )
54 self.base._build_model()
55 self.base._build_loss()
57 # datasets
58 self._build_model_assets()
59 self._build_data_transform()
60 self.base._build_dataset()
62 # dataloader
63 self._build_collate_fn()
64 self.base._build_dataloader()
66 # get_batch
67 self._build_get_batch()
68 self.base.attach_model_integration_data_pipeline()
69 self.base._compute_train_iters()
71 self.base._build_optimizer()
72 self.base.model_integration.attach_optimizer(self.base.optimizer)
73 self.base._build_lr_scheduler()
74 self.base._build_training_context()
75 self.base._build_graph_compiler()
76 self.base._init_callbacks()
78 def _build_model_assets(self) -> None:
79 """Build tokenizer-backed assets for text training."""
80 config: TrainerConfig = self.base.config
81 if config.dataset is None:
82 raise ValueError("dataset must define a build target")
83 assets_config = config.dataset.model_assets
84 tokenizer_target = assets_config.tokenizer
85 if tokenizer_target is None:
86 self.base.tokenizer = None
87 elif getattr(tokenizer_target, "pretrained_model_name_or_path", None):
88 self.base.tokenizer = tokenizer_target.build()
89 else:
90 self.base.tokenizer = tokenizer_target.build(
91 pretrained_model_name_or_path=config.model.pretrained_model_name_or_path,
92 )
94 self.base.chat_template = None
95 if assets_config.chat_template is not None:
96 if self.base.tokenizer is None:
97 raise ValueError("dataset.model_assets.tokenizer is required for conversation data")
98 if isinstance(assets_config.chat_template, str):
99 self.base.chat_template = build_chat_template(
100 assets_config.chat_template,
101 self.base.tokenizer,
102 )
103 else:
104 self.base.chat_template = assets_config.chat_template.build(
105 tokenizer=self.base.tokenizer,
106 )
108 self.base.model_assets = [self.base.model_config]
109 if self.base.tokenizer is not None:
110 self.base.model_assets.append(self.base.tokenizer)
111 if self.base.chat_template is not None:
112 self.base.model_assets.append(self.base.chat_template)
114 def _build_data_transform(self) -> None:
115 """Build the configured text sample transform."""
116 dataset_config = self.base.config.dataset
117 if dataset_config is None:
118 raise ValueError("dataset must define a build target")
119 if dataset_config.data_transform is None:
120 self.base.data_transform = None
121 return
122 self.base.data_transform = dataset_config.data_transform.build(
123 tokenizer=self.base.tokenizer,
124 chat_template=self.base.chat_template,
125 )
127 def _build_collate_fn(self) -> None:
128 """Build the text collator and gradient-accumulation batch count."""
129 dataloader_config = self.base.config.dataloader
130 if dataloader_config is None or dataloader_config.collate_fn is None:
131 raise ValueError("dataloader.collate_fn must define a build target")
132 training_config = self.base.config.training
133 self.base.num_micro_batches = calculate_num_micro_batches(
134 global_batch_size=training_config.global_batch_size,
135 micro_batch_size=training_config.micro_batch_size,
136 dp_world_size=self.base.mesh.dp_size,
137 )
138 self.base.collate_fn = dataloader_config.collate_fn.build(
139 mesh_context=self.base.mesh,
140 )
142 def _build_get_batch(self) -> None:
143 """Build the DataLoader-to-LLM batch runtime."""
144 config = self.base.config
145 if config.dataloader.get_batch is None:
146 raise ValueError("dataloader.get_batch must define a batching runtime target")
147 get_batch = config.dataloader.get_batch.build(
148 mesh_context=self.base.mesh,
149 device=self.base.device,
150 tokenizer=self.base.tokenizer,
151 data_config=getattr(config.dataset, "data_config", {}),
152 pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
153 )
154 self.base.get_batch = get_batch
156 @property
157 def distributed_setup(self) -> Any:
158 """Return the shared distributed setup."""
159 return self.base.distributed_setup
161 @property
162 def mesh(self) -> Any:
163 """Return the shared mesh context."""
164 return self.base.mesh
166 @property
167 def dp_cp_mesh(self) -> Any:
168 """Return the shared data/context-parallel mesh."""
169 return self.base.dp_cp_mesh
171 def on_train_begin(self) -> None:
172 """Dispatch the training-begin lifecycle hook."""
173 self.base.on_train_begin()
175 def on_train_end(self) -> None:
176 """Dispatch the training-end lifecycle hook."""
177 self.base.on_train_end()
179 def on_epoch_begin(self) -> None:
180 """Dispatch the epoch-begin lifecycle hook."""
181 self.base.on_epoch_begin()
183 def on_epoch_end(self) -> None:
184 """Dispatch the epoch-end lifecycle hook."""
185 self.base.on_epoch_end()
187 def on_step_begin(self) -> None:
188 """Dispatch the step-begin lifecycle hook."""
189 self.base.on_step_begin()
191 def on_micro_step_begin(self, micro_batch: dict[str, Any]) -> None:
192 """Dispatch the micro-step-begin lifecycle hook.
194 Args:
195 micro_batch: Prepared inputs and lightweight metric metadata.
196 """
197 self.base.on_micro_step_begin(micro_batch)
199 def on_step_end(
200 self,
201 loss: Any = None,
202 loss_dict: Any = None,
203 grad_norm: Any = None,
204 ) -> None:
205 """Dispatch the step-end lifecycle hook.
207 Args:
208 loss: Aggregated loss for the optimizer step.
209 loss_dict: Named loss values for the optimizer step.
210 grad_norm: Gradient norm measured before the optimizer update.
211 """
212 self.base.on_step_end(
213 loss=loss,
214 loss_dict=loss_dict,
215 grad_norm=grad_norm,
216 )
218 def forward_backward_step(
219 self,
220 data_iterator: Any,
221 num_micro_steps: int,
222 ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
223 """Fetch and execute one forward-backward micro-step.
225 Args:
226 data_iterator: Iterator providing one raw FB batch per call.
227 num_micro_steps: Number of FB steps in the current optimizer step.
229 Returns:
230 Loss tensor and named loss tensors for this FB step.
231 """
232 model_inputs, loss_inputs = self.base.get_batch(data_iterator)
233 self.base.current_token_counts = count_loss_token(loss_inputs)
234 self.base.step_token_counts = {
235 name: token_count * num_micro_steps
236 for name, token_count in self.base.current_token_counts.items()
237 }
238 metric_inputs = {
239 **model_inputs,
240 **loss_inputs,
241 "token_count": self.base.current_token_counts["foundation_tokens"],
242 }
243 self.on_micro_step_begin(metric_inputs)
244 del metric_inputs
245 loss, loss_dict = self.base.forward_backward_step(model_inputs, loss_inputs)
247 return loss, loss_dict
249 def train_step(self, data_iterator: Any) -> Dict[str, float]:
250 """Execute one text training step.
252 Args:
253 data_iterator: Iterator providing the step's micro-batches.
255 Returns:
256 Aggregated loss and gradient norm for the completed step.
257 """
258 memory_profiler.step()
259 num_micro_steps = self.base.num_micro_batches
261 self.on_step_begin()
262 self.base.model_integration.begin_step(self.base.state.global_step + 1)
263 synchronize()
265 total_loss = 0.0
266 total_loss_dict = defaultdict(int)
268 for micro_step in range(num_micro_steps):
269 self.base.model_reshard(micro_step, num_micro_steps)
270 self.base.configure_fsdp_gradient_sync(
271 micro_step,
272 num_micro_steps,
273 )
274 self.base.begin_fsdp_runtime_diagnostics(micro_step)
275 loss, loss_dict = self.forward_backward_step(
276 data_iterator,
277 num_micro_steps,
278 )
280 total_loss += loss.item()
281 for loss_name, loss_value in loss_dict.items():
282 total_loss_dict[loss_name] += loss_value.item()
284 grad_norm = self.base.prepare_optimizer_step()
285 self.base.step_optimizers_and_schedulers()
287 # Checkpoint and logging callbacks observe the number of completed
288 # optimizer updates.
289 self.base.state.global_step += 1
290 grad_norm_value = float(grad_norm)
291 self.base._end_model_integration_step( # pylint: disable=protected-access
292 {"loss": total_loss, "grad_norm": grad_norm_value}
293 )
294 self.on_step_end(
295 loss=total_loss,
296 loss_dict=total_loss_dict,
297 grad_norm=grad_norm_value,
298 )
300 return {
301 "loss": total_loss,
302 "grad_norm": grad_norm_value,
303 }
305 def train(self) -> None:
306 """Run the configured global optimizer steps by Dataset epoch."""
307 try:
308 config = self.base.config
309 self.on_train_begin()
310 logger.info(
311 "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
312 self.base.local_rank,
313 self.base.state.global_step,
314 self.base.train_iters,
315 self.base.state.epoch,
316 self.base.train_epochs,
317 )
319 # Checkpoint resume restores state.global_step, state.epoch, and the DataLoader cursor.
320 for epoch in range(self.base.state.epoch, self.base.train_epochs):
321 train_dataloader = self.base.train_dataloader
322 if hasattr(train_dataloader, "set_epoch"):
323 train_dataloader.set_epoch(epoch)
325 self.base.state.epoch = epoch
326 self.on_epoch_begin()
327 data_iterator = iter(train_dataloader) if train_dataloader is not None else None
328 start_step = self.base.state.global_step - epoch * self.base.train_steps
329 train_steps = min(self.base.train_steps, self.base.train_iters - epoch * self.base.train_steps)
330 for _ in range(start_step, train_steps):
331 try:
332 self.train_step(data_iterator)
333 except StopIteration:
334 logger.info(
335 "epoch:%s Dataloader finished with drop_last %s", epoch, config.dataloader.drop_last
336 )
337 break
339 self.on_epoch_end()
340 self.base.state.epoch = epoch + 1
341 print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")
342 if self.base.state.global_step >= self.base.train_iters:
343 break
345 memory_profiler.stop()
346 self.on_train_end()
348 except BaseException:
349 try:
350 memory_profiler.stop()
351 # Snapshot export must preserve the training error even if interrupted.
352 except BaseException: # pylint: disable=broad-exception-caught
353 logger.exception("Failed to save memory snapshot after training error.")
354 finally:
355 self.base.abort_profilers()
356 raise
358 synchronize()
359 self.base.destroy_distributed()