Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / vlm_trainer.py: 52%
199 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"""VLM Trainer assembled from the shared BaseTrainer stages."""
17__all__ = ["VLMTrainer"]
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.batching.get_batch import OmniParallelBatch
27from hyper_parallel.data.omni import OmniDataTransform
28from hyper_parallel.data.omni.metadata import vlm_sample_metadata
29from hyper_parallel.distributed.encoder_parallel import build_encoder_parallel_runtime
30from hyper_parallel.distributed.encoder_parallel.transport import EncoderParallelBatchTransport
31from hyper_parallel.models.registry import resolve_model_adapter
32from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
33from hyper_parallel.trainer.runtime.logging import create_logger
34from hyper_parallel.trainer.runtime.memory import print_device_mem_info
35from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
36from hyper_parallel.trainer.runtime.device import synchronize # pylint: disable=syntax-error
37from hyper_parallel.trainer.base import BaseTrainer
38from hyper_parallel.distributed_data import DistributedDataLoader
39from hyper_parallel.trainer.config import TrainerConfig
41logger = create_logger(__name__)
44class VLMTrainer:
45 """Compose the VLM training runtime from explicit BaseTrainer stages."""
47 base: BaseTrainer
49 def __init__(self, config: TrainerConfig) -> None:
50 """Build the VLM Trainer in the same explicit order as VeOmni."""
51 self.base = BaseTrainer.__new__(BaseTrainer)
52 self.base.config = config
54 self.base._setup()
55 memory_profiler.reset(
56 self.base.config.memory,
57 global_rank=self.base.global_rank,
58 tp_rank=self.base.mesh.tp_rank,
59 dp_rank=self.base.mesh.dp_rank,
60 )
61 self.base._build_model()
62 self.base._build_loss()
64 # datasets
65 self._build_model_assets()
66 self._build_data_transform()
67 self.base._build_dataset()
69 # dataloader
70 self._build_collate_fn()
71 self.base._build_dataloader(metadata_fn=vlm_sample_metadata)
73 # get_batch
74 self._build_get_batch()
75 self.base.attach_model_integration_data_pipeline()
76 self.base._compute_train_iters()
78 self.base._build_optimizer()
79 self.base.model_integration.attach_optimizer(self.base.optimizer)
80 self.base._build_lr_scheduler()
81 self.base._build_training_context()
82 self.base._init_callbacks()
84 def _build_model_assets(self) -> None:
85 """Build and publish the processor consumed by the Omni transform."""
86 config: TrainerConfig = self.base.config
87 if config.dataset is None:
88 raise ValueError("dataset must define a build target")
90 # The model-owned processor is built through the Omni ``model_assets``
91 # target; forward the model's ``trust_remote_code`` flag so a natively
92 # registered architecture such as the Kimi-K2.6 family (``kimi_k26``)
93 # stays on the local implementation unless the config opts into remote
94 # code.
95 processor = config.dataset.model_assets.build(
96 trust_remote_code=bool(
97 getattr(config.model, "trust_remote_code", True)
98 ),
99 )
100 tokenizer = getattr(processor, "tokenizer", None)
101 if tokenizer is None:
102 raise ValueError("dataset.model_assets must build a processor with a tokenizer")
104 self.base.processor = processor
105 self.base.tokenizer = tokenizer
106 self.base.chat_template = getattr(processor, "chat_template", None)
107 self.base.model_assets = [self.base.model_config, processor]
109 def _build_data_transform(self) -> None:
110 """Build the Omni encoding lifecycle from the prebuilt processor."""
111 dataset_config = self.base.config.dataset
112 if dataset_config is None:
113 raise ValueError("dataset must define a build target")
115 if dataset_config.data_transform is None:
116 raise ValueError("dataset.data_transform must define an Omni transform build target")
118 if dataset_config.training_mode is None:
119 raise ValueError("dataset.training_mode must be explicitly set to 'sft' or 'pretrain' for VLM training")
121 data_transform = dataset_config.data_transform.build(
122 processor=self.base.processor,
123 training_mode=dataset_config.training_mode,
124 )
125 if not isinstance(data_transform, OmniDataTransform):
126 raise TypeError("dataset.data_transform must build an OmniDataTransform")
128 self.base.data_transform = data_transform
130 def _build_collate_fn(self) -> None:
131 """Build the VLM collator and gradient-accumulation batch count."""
132 dataloader_config = self.base.config.dataloader
133 if dataloader_config is None or dataloader_config.collate_fn is None:
134 raise ValueError("dataloader.collate_fn must define a build target")
135 training_config = self.base.config.training
136 self.base.num_micro_batches = calculate_num_micro_batches(
137 global_batch_size=training_config.global_batch_size,
138 micro_batch_size=training_config.micro_batch_size,
139 dp_world_size=self.base.mesh.dp_size,
140 )
141 self.base.collate_fn = dataloader_config.collate_fn.build(
142 mesh_context=self.base.mesh,
143 )
145 def _build_get_batch(self) -> None:
146 """Build the DataLoader-to-LLM batch runtime."""
147 config = self.base.config
148 if config.dataloader.get_batch is None:
149 raise ValueError("dataloader.get_batch must define a batching runtime target")
150 if getattr(config.dataloader.get_batch, "encoder_parallel", None) is not None:
151 raise ValueError("encoder_parallel belongs under accelerator, not dataloader.get_batch")
152 encoder_runtime = build_encoder_parallel_runtime(
153 config=config.accelerator.encoder_parallel,
154 mesh_context=self.base.mesh,
155 model=self.base.model,
156 model_adapter_spec=resolve_model_adapter(self.base.model),
157 execution_options={
158 "micro_batch_size": int(config.training.micro_batch_size),
159 "compile_enabled": bool(config.compile.enabled),
160 "validate_placement": bool(getattr(config.model, "validate_placement", False)),
161 },
162 )
163 get_batch = config.dataloader.get_batch.build(
164 mesh_context=self.base.mesh,
165 dataloader_context=self.base.dataloader_context,
166 device=self.base.device,
167 data_config=getattr(config.dataset, "data_config", {}),
168 pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
169 )
170 if encoder_runtime is not None:
171 if not isinstance(get_batch, OmniParallelBatch):
172 raise ValueError("encoder parallel requires OmniParallelBatch")
173 get_batch.batch_transport = EncoderParallelBatchTransport(runtime=encoder_runtime)
174 self.base.get_batch = get_batch
176 @property
177 def distributed_setup(self) -> Any:
178 """Return the shared distributed setup."""
179 return self.base.distributed_setup
181 @property
182 def mesh(self) -> Any:
183 """Return the shared mesh context."""
184 return self.base.mesh
186 @property
187 def dp_cp_mesh(self) -> Any:
188 """Return the shared data/context-parallel mesh."""
189 return self.base.dp_cp_mesh
191 def on_train_begin(self) -> None:
192 """Dispatch the training-begin lifecycle hook."""
193 self.base.on_train_begin()
195 def on_train_end(self) -> None:
196 """Dispatch the training-end lifecycle hook."""
197 self.base.on_train_end()
199 def on_epoch_begin(self) -> None:
200 """Dispatch the epoch-begin lifecycle hook."""
201 self.base.on_epoch_begin()
203 def on_epoch_end(self) -> None:
204 """Dispatch the epoch-end lifecycle hook."""
205 self.base.on_epoch_end()
207 def on_step_begin(self) -> None:
208 """Dispatch the step-begin lifecycle hook."""
209 self.base.on_step_begin()
211 def on_micro_step_begin(self, micro_batch: dict[str, Any]) -> None:
212 """Dispatch the micro-step-begin lifecycle hook.
214 Args:
215 micro_batch: Prepared inputs and lightweight metric metadata.
216 """
217 self.base.on_micro_step_begin(micro_batch)
219 def on_micro_step_end(self, loss_dict: dict[str, torch.Tensor]) -> None:
220 """Dispatch the micro-step-end lifecycle hook.
222 Args:
223 loss_dict: Named losses and metrics from the completed forward-backward step.
224 """
225 self.base.on_micro_step_end(loss_dict)
227 def on_step_end(
228 self,
229 loss: Any = None,
230 loss_dict: Any = None,
231 grad_norm: Any = None,
232 ) -> None:
233 """Dispatch the step-end lifecycle hook.
235 Args:
236 loss: Aggregate step loss.
237 loss_dict: Named step losses.
238 grad_norm: Measured gradient norm.
239 """
240 self.base.on_step_end(
241 loss=loss,
242 loss_dict=loss_dict,
243 grad_norm=grad_norm,
244 )
246 def forward_backward_step(
247 self,
248 data_iterator: Any,
249 ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
250 """Fetch and execute one forward-backward micro-step.
252 Args:
253 data_iterator: Iterator providing one multimodal batch per call.
255 Returns:
256 Loss tensor and named loss tensors for this FB step.
257 """
258 model_inputs, loss_inputs = self.base.get_batch(data_iterator)
259 self.base.current_token_counts = count_loss_token(loss_inputs)
260 self.on_micro_step_begin(model_inputs)
261 loss, loss_dict = self.base.forward_backward_step(model_inputs, loss_inputs)
262 self.on_micro_step_end(loss_dict)
264 return loss, loss_dict
266 def train_step(self, data_iterator: Any) -> Dict[str, float]:
267 """Execute one VLM training step.
269 Args:
270 data_iterator: Iterator providing multimodal micro-batches.
272 Returns:
273 Aggregated loss and gradient norm.
274 """
275 memory_profiler.step()
276 num_micro_steps = self.base.num_micro_batches
277 self.base.begin_loss_step()
279 self.on_step_begin()
280 self.base.model_integration.begin_step(self.base.state.global_step + 1)
281 synchronize()
283 total_loss = 0.0
284 total_loss_dict = defaultdict(int)
286 for micro_step in range(num_micro_steps):
287 self.base.model_reshard(micro_step, num_micro_steps)
288 self.base.configure_fsdp_gradient_sync(
289 micro_step,
290 num_micro_steps,
291 )
292 self.base.begin_fsdp_runtime_diagnostics(micro_step)
293 loss, loss_dict = self.forward_backward_step(
294 data_iterator,
295 )
297 total_loss += loss.item()
298 for loss_name, loss_value in loss_dict.items():
299 total_loss_dict[loss_name] += loss_value.item()
301 total_loss, grad_norm = self.base.prepare_optimizer_step(total_loss, total_loss_dict)
302 self.base.step_optimizers_and_schedulers()
304 # Checkpoint and logging callbacks observe completed optimizer updates.
305 self.base.state.global_step += 1
306 grad_norm_value = float(grad_norm)
307 # The base trainer owns this integration hook and the VLM subclass only
308 # mirrors the base step order, so reaching the private method is intended.
309 self.base._end_model_integration_step( # pylint: disable=protected-access
310 {"loss": total_loss, "grad_norm": grad_norm_value}
311 )
312 self.on_step_end(
313 loss=total_loss,
314 loss_dict=total_loss_dict,
315 grad_norm=grad_norm_value,
316 )
318 step_metrics = {
319 "loss": total_loss,
320 "grad_norm": grad_norm_value,
321 }
322 return step_metrics
324 def _train_epoch(self, epoch: int, balanced_loader: bool) -> bool:
325 """Run complete optimizer steps and report actual Dataset exhaustion."""
326 train_dataloader = self.base.train_dataloader
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 if balanced_loader:
331 start_step = 0
332 train_steps = self.base.train_iters - self.base.state.global_step
333 for _ in range(start_step, train_steps):
334 try:
335 self.train_step(data_iterator)
336 except StopIteration:
337 optimizers = self.base.optimizer
338 for optimizer in optimizers if isinstance(optimizers, list) else [optimizers]:
339 optimizer.zero_grad()
340 logger.info(
341 "epoch:%s Dataloader finished with drop_last %s", epoch, self.base.config.dataloader.drop_last
342 )
343 return True
344 return False
346 def train(self) -> None:
347 """Run the VLM training loop."""
348 try:
349 self.on_train_begin()
350 logger.info(
351 "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
352 self.base.local_rank,
353 self.base.state.global_step,
354 self.base.train_iters,
355 self.base.state.epoch,
356 self.base.train_epochs,
357 )
359 train_dataloader = self.base.train_dataloader
360 balanced_loader = isinstance(train_dataloader, DistributedDataLoader)
361 start_epoch = self.base.state.epoch
362 epochs = count(start_epoch) if balanced_loader else range(start_epoch, self.base.train_epochs)
363 for epoch in epochs:
364 if self.base.state.global_step >= self.base.train_iters:
365 break
366 # Do not erase the cursor restored by on_train_begin on the first epoch.
367 if hasattr(train_dataloader, "set_epoch") and (not balanced_loader or epoch != start_epoch):
368 train_dataloader.set_epoch(epoch)
370 self.base.state.epoch = epoch
371 self.on_epoch_begin()
372 epoch_start_step = self.base.state.global_step
373 epoch_exhausted = self._train_epoch(epoch, balanced_loader)
374 if balanced_loader and epoch != start_epoch and self.base.state.global_step == epoch_start_step:
375 raise ValueError("VLM Dataset epoch cannot provide one complete optimizer step")
377 self.on_epoch_end()
378 self.base.state.epoch = epoch + 1 if not balanced_loader or epoch_exhausted else epoch
379 print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")
381 memory_profiler.stop()
382 self.on_train_end()
384 except BaseException:
385 try:
386 memory_profiler.stop()
387 # Snapshot export must preserve the training error even if interrupted.
388 except BaseException: # pylint: disable=broad-exception-caught
389 logger.exception("Failed to save memory snapshot after training error.")
390 finally:
391 self.base.abort_profilers()
392 raise
394 if balanced_loader:
395 try:
396 train_dataloader.wait_for_prefetch()
397 finally:
398 train_dataloader.close()
400 synchronize()
401 self.base.destroy_distributed()