Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / text_trainer.py: 46%
201 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.
15"""Text Trainer assembled from the shared BaseTrainer stages."""
17__all__ = ["TextTrainer"]
19from collections import defaultdict
20from itertools import count
21from typing import Any, Dict
23import torch # pylint: disable=forbidden-backend-import
25from hyper_parallel.data.batching import calculate_num_micro_batches
26from hyper_parallel.data.text import build_chat_template
27from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
28from hyper_parallel.trainer.runtime.logging import create_logger
29from hyper_parallel.trainer.runtime.memory import print_device_mem_info
30from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
31from hyper_parallel.trainer.runtime.device import synchronize
32from hyper_parallel.trainer.base import BaseTrainer
33from hyper_parallel.distributed_data import DistributedDataLoader
34from hyper_parallel.trainer.config import TrainerConfig
36logger = create_logger(__name__)
39class TextTrainer:
40 """Compose the text training runtime from explicit BaseTrainer stages."""
42 base: BaseTrainer
44 def __init__(self, config: TrainerConfig) -> None:
45 """Build the text Trainer in the same explicit order as VeOmni."""
46 self.base = BaseTrainer.__new__(BaseTrainer)
47 self.base.config = config
49 self.base._setup()
50 memory_profiler.reset(
51 self.base.config.memory,
52 global_rank=self.base.global_rank,
53 tp_rank=self.base.mesh.tp_rank,
54 dp_rank=self.base.mesh.dp_rank,
55 )
56 self.base._build_model()
57 self.base._build_loss()
59 # datasets
60 self._build_model_assets()
61 self._build_data_transform()
62 self.base._build_dataset()
64 # dataloader
65 self._build_collate_fn()
66 self.base._build_dataloader()
68 # get_batch
69 self._build_get_batch()
70 self.base.attach_model_integration_data_pipeline()
71 self.base._compute_train_iters()
73 self.base._build_optimizer()
74 self.base.model_integration.attach_optimizer(self.base.optimizer)
75 self.base._build_lr_scheduler()
76 self.base._build_training_context()
77 self.base._build_graph_compiler()
78 self.base._init_callbacks()
80 def _build_model_assets(self) -> None:
81 """Build tokenizer-backed assets for text training."""
82 config: TrainerConfig = self.base.config
83 if config.dataset is None:
84 raise ValueError("dataset must define a build target")
85 assets_config = config.dataset.model_assets
86 tokenizer_target = assets_config.tokenizer
87 if tokenizer_target is None:
88 self.base.tokenizer = None
89 elif getattr(tokenizer_target, "pretrained_model_name_or_path", None):
90 self.base.tokenizer = tokenizer_target.build()
91 else:
92 self.base.tokenizer = tokenizer_target.build(
93 pretrained_model_name_or_path=config.model.pretrained_model_name_or_path,
94 )
96 self.base.chat_template = None
97 if assets_config.chat_template is not None:
98 if self.base.tokenizer is None:
99 raise ValueError("dataset.model_assets.tokenizer is required for conversation data")
100 if isinstance(assets_config.chat_template, str):
101 self.base.chat_template = build_chat_template(
102 assets_config.chat_template,
103 self.base.tokenizer,
104 chat_template_kwargs=assets_config.chat_template_kwargs,
105 log_first_rendered_template=config.debug.log_first_rendered_template,
106 )
107 else:
108 if assets_config.chat_template_kwargs or config.debug.log_first_rendered_template:
109 raise ValueError(
110 "dataset.model_assets chat-template options require chat_template='tokenizer'"
111 )
112 self.base.chat_template = assets_config.chat_template.build(
113 tokenizer=self.base.tokenizer,
114 )
116 self.base.model_assets = [self.base.model_config]
117 if self.base.tokenizer is not None:
118 self.base.model_assets.append(self.base.tokenizer)
119 if self.base.chat_template is not None:
120 self.base.model_assets.append(self.base.chat_template)
122 def _build_data_transform(self) -> None:
123 """Build the configured text sample transform."""
124 dataset_config = self.base.config.dataset
125 if dataset_config is None:
126 raise ValueError("dataset must define a build target")
127 if dataset_config.data_transform is None:
128 self.base.data_transform = None
129 return
130 self.base.data_transform = dataset_config.data_transform.build(
131 tokenizer=self.base.tokenizer,
132 chat_template=self.base.chat_template,
133 )
135 def _build_collate_fn(self) -> None:
136 """Build the text collator and gradient-accumulation batch count."""
137 dataloader_config = self.base.config.dataloader
138 if dataloader_config is None or dataloader_config.collate_fn is None:
139 raise ValueError("dataloader.collate_fn must define a build target")
140 training_config = self.base.config.training
141 self.base.num_micro_batches = calculate_num_micro_batches(
142 global_batch_size=training_config.global_batch_size,
143 micro_batch_size=training_config.micro_batch_size,
144 dp_world_size=self.base.mesh.dp_size,
145 )
146 self.base.collate_fn = dataloader_config.collate_fn.build(
147 mesh_context=self.base.mesh,
148 )
150 def _build_get_batch(self) -> None:
151 """Build the DataLoader-to-LLM batch runtime."""
152 config = self.base.config
153 if config.dataloader.get_batch is None:
154 raise ValueError("dataloader.get_batch must define a batching runtime target")
155 get_batch = config.dataloader.get_batch.build(
156 mesh_context=self.base.mesh,
157 dataloader_context=self.base.dataloader_context,
158 device=self.base.device,
159 tokenizer=self.base.tokenizer,
160 data_config=getattr(config.dataset, "data_config", {}),
161 pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
162 )
163 self.base.get_batch = get_batch
165 @property
166 def distributed_setup(self) -> Any:
167 """Return the shared distributed setup."""
168 return self.base.distributed_setup
170 @property
171 def mesh(self) -> Any:
172 """Return the shared mesh context."""
173 return self.base.mesh
175 @property
176 def dp_cp_mesh(self) -> Any:
177 """Return the shared data/context-parallel mesh."""
178 return self.base.dp_cp_mesh
180 def on_train_begin(self) -> None:
181 """Dispatch the training-begin lifecycle hook."""
182 self.base.on_train_begin()
184 def on_train_end(self) -> None:
185 """Dispatch the training-end lifecycle hook."""
186 self.base.on_train_end()
188 def on_epoch_begin(self) -> None:
189 """Dispatch the epoch-begin lifecycle hook."""
190 self.base.on_epoch_begin()
192 def on_epoch_end(self) -> None:
193 """Dispatch the epoch-end lifecycle hook."""
194 self.base.on_epoch_end()
196 def on_step_begin(self) -> None:
197 """Dispatch the step-begin lifecycle hook."""
198 self.base.on_step_begin()
200 def on_micro_step_begin(self, micro_batch: dict[str, Any]) -> None:
201 """Dispatch the micro-step-begin lifecycle hook.
203 Args:
204 micro_batch: Prepared inputs and lightweight metric metadata.
205 """
206 self.base.on_micro_step_begin(micro_batch)
208 def on_micro_step_end(self, loss_dict: dict[str, torch.Tensor]) -> None:
209 """Dispatch the micro-step-end lifecycle hook.
211 Args:
212 loss_dict: Named losses and metrics from the completed forward-backward step.
213 """
214 self.base.on_micro_step_end(loss_dict)
216 def on_step_end(
217 self,
218 loss: Any = None,
219 loss_dict: Any = None,
220 grad_norm: Any = None,
221 ) -> None:
222 """Dispatch the step-end lifecycle hook.
224 Args:
225 loss: Aggregated loss for the optimizer step.
226 loss_dict: Named loss values for the optimizer step.
227 grad_norm: Gradient norm measured before the optimizer update.
228 """
229 self.base.on_step_end(
230 loss=loss,
231 loss_dict=loss_dict,
232 grad_norm=grad_norm,
233 )
235 def forward_backward_step(
236 self,
237 data_iterator: Any,
238 ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
239 """Fetch and execute one forward-backward micro-step.
241 Args:
242 data_iterator: Iterator providing one raw FB batch per call.
244 Returns:
245 Loss tensor and named loss tensors for this FB step.
246 """
247 model_inputs, loss_inputs = self.base.get_batch(data_iterator)
248 self.base.current_token_counts = count_loss_token(loss_inputs)
249 metric_inputs = {
250 **model_inputs,
251 **loss_inputs,
252 "token_count": self.base.current_token_counts["foundation_tokens"],
253 }
254 self.on_micro_step_begin(metric_inputs)
255 del metric_inputs
256 loss, loss_dict = self.base.forward_backward_step(model_inputs, loss_inputs)
257 self.on_micro_step_end(loss_dict)
259 return loss, loss_dict
261 def train_step(self, data_iterator: Any) -> Dict[str, float]:
262 """Execute one text training step.
264 Args:
265 data_iterator: Iterator providing the step's micro-batches.
267 Returns:
268 Aggregated loss and gradient norm for the completed step.
269 """
270 memory_profiler.step()
271 num_micro_steps = self.base.num_micro_batches
272 self.base.begin_loss_step()
274 self.on_step_begin()
275 self.base.model_integration.begin_step(self.base.state.global_step + 1)
276 synchronize()
278 total_loss = 0.0
279 total_loss_dict = defaultdict(int)
281 for micro_step in range(num_micro_steps):
282 self.base.model_reshard(micro_step, num_micro_steps)
283 self.base.configure_fsdp_gradient_sync(
284 micro_step,
285 num_micro_steps,
286 )
287 self.base.begin_fsdp_runtime_diagnostics(micro_step)
288 loss, loss_dict = self.forward_backward_step(
289 data_iterator,
290 )
292 total_loss += loss.item()
293 for loss_name, loss_value in loss_dict.items():
294 total_loss_dict[loss_name] += loss_value.item()
296 total_loss, grad_norm = self.base.prepare_optimizer_step(total_loss, total_loss_dict)
297 self.base.step_optimizers_and_schedulers()
299 # Checkpoint and logging callbacks observe the number of completed
300 # optimizer updates.
301 self.base.state.global_step += 1
302 grad_norm_value = float(grad_norm)
303 self.base._end_model_integration_step( # pylint: disable=protected-access
304 {"loss": total_loss, "grad_norm": grad_norm_value}
305 )
306 self.on_step_end(
307 loss=total_loss,
308 loss_dict=total_loss_dict,
309 grad_norm=grad_norm_value,
310 )
312 return {
313 "loss": total_loss,
314 "grad_norm": grad_norm_value,
315 }
317 def _train_epoch(self, epoch: int, balanced_loader: bool) -> bool:
318 """Consume complete optimizer steps and report whether the source epoch ended."""
319 train_dataloader = self.base.train_dataloader
320 data_iterator = iter(train_dataloader) if train_dataloader is not None else None
321 start_step = self.base.state.global_step - epoch * self.base.train_steps
322 train_steps = min(self.base.train_steps, self.base.train_iters - epoch * self.base.train_steps)
323 if balanced_loader:
324 start_step = 0
325 train_steps = self.base.train_iters - self.base.state.global_step
326 for _ in range(start_step, train_steps):
327 try:
328 self.train_step(data_iterator)
329 except StopIteration:
330 if balanced_loader:
331 optimizers = self.base.optimizer
332 for optimizer in optimizers if isinstance(optimizers, list) else [optimizers]:
333 optimizer.zero_grad()
334 logger.info(
335 "epoch:%s Dataloader finished with drop_last %s", epoch, self.base.config.dataloader.drop_last
336 )
337 return True
338 return False
340 def train(self) -> None:
341 """Run the configured global optimizer steps by Dataset epoch."""
342 try:
343 self.on_train_begin()
344 logger.info(
345 "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
346 self.base.local_rank,
347 self.base.state.global_step,
348 self.base.train_iters,
349 self.base.state.epoch,
350 self.base.train_epochs,
351 )
353 train_dataloader = self.base.train_dataloader
354 balanced_loader = isinstance(train_dataloader, DistributedDataLoader)
355 start_epoch = self.base.state.epoch
356 epochs = count(start_epoch) if balanced_loader else range(start_epoch, self.base.train_epochs)
357 for epoch in epochs:
358 if self.base.state.global_step >= self.base.train_iters:
359 break
360 # The first collective iterator may hold a restored checkpoint cursor.
361 if hasattr(train_dataloader, "set_epoch") and (not balanced_loader or epoch != start_epoch):
362 train_dataloader.set_epoch(epoch)
364 self.base.state.epoch = epoch
365 self.on_epoch_begin()
366 epoch_start_step = self.base.state.global_step
367 epoch_exhausted = self._train_epoch(epoch, balanced_loader)
369 if balanced_loader and epoch != start_epoch and self.base.state.global_step == epoch_start_step:
370 raise ValueError("Indexed source epoch cannot provide one complete optimizer step")
372 self.on_epoch_end()
373 self.base.state.epoch = epoch + 1 if not balanced_loader or epoch_exhausted else epoch
374 print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")
375 if self.base.state.global_step >= self.base.train_iters:
376 break
378 memory_profiler.stop()
379 self.on_train_end()
381 except BaseException:
382 try:
383 memory_profiler.stop()
384 # Snapshot export must preserve the training error even if interrupted.
385 except BaseException: # pylint: disable=broad-exception-caught
386 logger.exception("Failed to save memory snapshot after training error.")
387 finally:
388 self.base.abort_profilers()
389 raise
391 if balanced_loader:
392 try:
393 train_dataloader.wait_for_prefetch()
394 finally:
395 train_dataloader.close()
397 synchronize()
398 self.base.destroy_distributed()