Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / vlm_trainer.py: 35%
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"""VLM Trainer assembled from the shared BaseTrainer stages."""
17__all__ = ["VLMTrainer"]
19from collections import defaultdict
20from typing import Any, Dict
22from hyper_parallel.data.batching import calculate_num_micro_batches
23from hyper_parallel.data.batching.get_batch import OmniParallelBatch
24from hyper_parallel.data.omni import OmniDataTransform
25from hyper_parallel.distributed.encoder_parallel import build_encoder_parallel_runtime
26from hyper_parallel.distributed.encoder_parallel.batch_input import EncoderParallelInputAdapter
27from hyper_parallel.models.registry import resolve_model_adapter
28from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
29from hyper_parallel.trainer.runtime.logging import create_logger
30from hyper_parallel.trainer.runtime.memory import print_device_mem_info
31from hyper_parallel.trainer.runtime.memory_profiler import memory_profiler
32from hyper_parallel.trainer.runtime.device import synchronize # pylint: disable=syntax-error
33from hyper_parallel.trainer.base import BaseTrainer
34from hyper_parallel.trainer.config import TrainerConfig
36logger = create_logger(__name__)
39class VLMTrainer:
40 """Compose the VLM training runtime from explicit BaseTrainer stages."""
42 base: BaseTrainer
44 def __init__(self, config: TrainerConfig) -> None:
45 """Build the VLM 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._init_callbacks()
79 def _build_model_assets(self) -> None:
80 """Build and publish the processor consumed by the Omni transform."""
81 config: TrainerConfig = self.base.config
82 if config.dataset is None:
83 raise ValueError("dataset must define a build target")
85 # The model-owned processor is built through the Omni ``model_assets``
86 # target; forward the model's ``trust_remote_code`` flag so a natively
87 # registered architecture such as the Kimi-K2.6 family (``kimi_k26``)
88 # stays on the local implementation unless the config opts into remote
89 # code.
90 processor = config.dataset.model_assets.build(
91 trust_remote_code=bool(
92 getattr(config.model, "trust_remote_code", True)
93 ),
94 )
95 tokenizer = getattr(processor, "tokenizer", None)
96 if tokenizer is None:
97 raise ValueError("dataset.model_assets must build a processor with a tokenizer")
99 self.base.processor = processor
100 self.base.tokenizer = tokenizer
101 self.base.chat_template = getattr(processor, "chat_template", None)
102 self.base.model_assets = [self.base.model_config, processor]
104 def _build_data_transform(self) -> None:
105 """Build the Omni encoding lifecycle from the prebuilt processor."""
106 dataset_config = self.base.config.dataset
107 if dataset_config is None:
108 raise ValueError("dataset must define a build target")
109 if dataset_config.data_transform is None:
110 raise ValueError("dataset.data_transform must define an Omni transform build target")
111 data_transform = dataset_config.data_transform.build(
112 processor=self.base.processor,
113 )
114 if not isinstance(data_transform, OmniDataTransform):
115 raise TypeError("dataset.data_transform must build an OmniDataTransform")
116 self.base.data_transform = data_transform
118 def _build_collate_fn(self) -> None:
119 """Build the VLM collator and gradient-accumulation batch count."""
120 dataloader_config = self.base.config.dataloader
121 if dataloader_config is None or dataloader_config.collate_fn is None:
122 raise ValueError("dataloader.collate_fn must define a build target")
123 training_config = self.base.config.training
124 self.base.num_micro_batches = calculate_num_micro_batches(
125 global_batch_size=training_config.global_batch_size,
126 micro_batch_size=training_config.micro_batch_size,
127 dp_world_size=self.base.mesh.dp_size,
128 )
129 self.base.collate_fn = dataloader_config.collate_fn.build(
130 mesh_context=self.base.mesh,
131 )
133 def _build_get_batch(self) -> None:
134 """Build the DataLoader-to-LLM batch runtime."""
135 config = self.base.config
136 if config.dataloader.get_batch is None:
137 raise ValueError("dataloader.get_batch must define a batching runtime target")
138 if getattr(config.dataloader.get_batch, "encoder_parallel", None) is not None:
139 raise ValueError("encoder_parallel belongs under accelerator, not dataloader.get_batch")
140 encoder_runtime = build_encoder_parallel_runtime(
141 config=config.accelerator.encoder_parallel,
142 mesh_context=self.base.mesh,
143 model=self.base.model,
144 model_adapter_spec=resolve_model_adapter(self.base.model),
145 execution_options={
146 "micro_batch_size": int(config.training.micro_batch_size),
147 "compile_enabled": bool(config.compile.enabled),
148 "validate_placement": bool(getattr(config.model, "validate_placement", False)),
149 },
150 )
151 get_batch = config.dataloader.get_batch.build(
152 mesh_context=self.base.mesh,
153 device=self.base.device,
154 data_config=getattr(config.dataset, "data_config", {}),
155 pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
156 )
157 if encoder_runtime is not None:
158 if not isinstance(get_batch, OmniParallelBatch):
159 raise ValueError("encoder parallel requires OmniParallelBatch")
160 get_batch.runtime_input_adapter = EncoderParallelInputAdapter(
161 runtime=encoder_runtime,
162 inner=get_batch.runtime_input_adapter,
163 )
164 self.base.get_batch = get_batch
166 @property
167 def distributed_setup(self) -> Any:
168 """Return the shared distributed setup."""
169 return self.base.distributed_setup
171 @property
172 def mesh(self) -> Any:
173 """Return the shared mesh context."""
174 return self.base.mesh
176 @property
177 def dp_cp_mesh(self) -> Any:
178 """Return the shared data/context-parallel mesh."""
179 return self.base.dp_cp_mesh
181 def on_train_begin(self) -> None:
182 """Dispatch the training-begin lifecycle hook."""
183 self.base.on_train_begin()
185 def on_train_end(self) -> None:
186 """Dispatch the training-end lifecycle hook."""
187 self.base.on_train_end()
189 def on_epoch_begin(self) -> None:
190 """Dispatch the epoch-begin lifecycle hook."""
191 self.base.on_epoch_begin()
193 def on_epoch_end(self) -> None:
194 """Dispatch the epoch-end lifecycle hook."""
195 self.base.on_epoch_end()
197 def on_step_begin(self, micro_batches: Any = None) -> None:
198 """Dispatch the step-begin lifecycle hook.
200 Args:
201 micro_batches: Prepared model inputs for this optimizer step.
202 """
203 self.base.on_step_begin(micro_batches=micro_batches)
205 def on_step_end(
206 self,
207 loss: Any = None,
208 loss_dict: Any = None,
209 grad_norm: Any = None,
210 ) -> None:
211 """Dispatch the step-end lifecycle hook.
213 Args:
214 loss: Aggregate step loss.
215 loss_dict: Named step losses.
216 grad_norm: Measured gradient norm.
217 """
218 self.base.on_step_end(
219 loss=loss,
220 loss_dict=loss_dict,
221 grad_norm=grad_norm,
222 )
224 def _forward_backward_micro_batches(
225 self,
226 training_batches: list[Any],
227 num_micro_steps: int,
228 ) -> tuple[float, Dict[str, float]]:
229 """Run and aggregate the VLM forward-backward micro-steps."""
230 total_loss = 0.0
231 total_loss_dict = defaultdict(int)
233 for micro_step, batch in enumerate(training_batches):
234 model_inputs, loss_inputs = batch
235 self.base.model_reshard(micro_step, num_micro_steps)
236 self.base.configure_fsdp_gradient_sync(
237 micro_step,
238 num_micro_steps,
239 )
240 self.base.begin_fsdp_runtime_diagnostics(micro_step)
241 self.base.current_token_counts = count_loss_token(loss_inputs)
242 self.base.step_token_counts = {
243 name: token_count * num_micro_steps
244 for name, token_count in self.base.current_token_counts.items()
245 }
246 loss, loss_dict = self.base.forward_backward_step(model_inputs, loss_inputs)
248 # Release each device batch as soon as its backward pass completes;
249 # prefetching all micro-batches must not pin them until step end.
250 training_batches[micro_step] = None
252 total_loss += loss.item()
253 for loss_name, loss_value in loss_dict.items():
254 total_loss_dict[loss_name] += loss_value.item()
256 return total_loss, total_loss_dict
258 def train_step(self, data_iterator: Any) -> Dict[str, float]:
259 """Execute one VLM training step.
261 Args:
262 data_iterator: Iterator providing multimodal micro-batches.
264 Returns:
265 Aggregated loss and gradient norm.
266 """
267 memory_profiler.step()
268 first_training_batch = self.base.get_batch(data_iterator)
269 num_micro_steps = self.base.num_micro_batches
270 training_batches = [first_training_batch]
271 for _ in range(1, num_micro_steps):
272 training_batches.append(self.base.get_batch(data_iterator))
274 self.on_step_begin(
275 micro_batches=[model_inputs for model_inputs, _ in training_batches]
276 )
277 self.base.model_integration.begin_step(self.base.state.global_step + 1)
278 synchronize()
280 total_loss, total_loss_dict = self._forward_backward_micro_batches(
281 training_batches,
282 num_micro_steps,
283 )
285 grad_norm = self.base.prepare_optimizer_step()
286 self.base.step_optimizers_and_schedulers()
288 # Checkpoint and logging callbacks observe completed optimizer updates.
289 self.base.state.global_step += 1
290 grad_norm_value = float(grad_norm)
291 # The base trainer owns this integration hook and the VLM subclass only
292 # mirrors the base step order, so reaching the private method is intended.
293 self.base._end_model_integration_step( # pylint: disable=protected-access
294 {"loss": total_loss, "grad_norm": grad_norm_value}
295 )
296 self.on_step_end(
297 loss=total_loss,
298 loss_dict=total_loss_dict,
299 grad_norm=grad_norm_value,
300 )
302 step_metrics = {
303 "loss": total_loss,
304 "grad_norm": grad_norm_value,
305 }
306 return step_metrics
308 def train(self) -> None:
309 """Run the VLM training loop."""
310 try:
311 config = self.base.config
312 self.on_train_begin()
313 logger.info(
314 "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
315 self.base.local_rank,
316 self.base.state.global_step,
317 self.base.train_iters,
318 self.base.state.epoch,
319 self.base.train_epochs,
320 )
322 # Checkpoint resume restores state.global_step, state.epoch, and the DataLoader cursor.
323 for epoch in range(self.base.state.epoch, self.base.train_epochs):
324 train_dataloader = self.base.train_dataloader
326 if hasattr(train_dataloader, "set_epoch"):
327 train_dataloader.set_epoch(epoch)
329 self.on_epoch_begin()
330 data_iterator = iter(train_dataloader) if train_dataloader is not None else None
332 start_step = self.base.state.global_step - epoch * self.base.train_steps
333 train_steps = min(self.base.train_steps, self.base.train_iters - epoch * self.base.train_steps)
334 for _ in range(start_step, train_steps):
335 try:
336 self.train_step(data_iterator)
337 except StopIteration:
338 logger.info(
339 "epoch:%s Dataloader finished with drop_last %s", epoch, config.dataloader.drop_last
340 )
341 break
343 self.on_epoch_end()
344 self.base.state.epoch = epoch + 1
345 print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")
347 if self.base.state.global_step >= self.base.train_iters:
348 break
350 memory_profiler.stop()
351 self.on_train_end()
353 except BaseException:
354 try:
355 memory_profiler.stop()
356 # Snapshot export must preserve the training error even if interrupted.
357 except BaseException: # pylint: disable=broad-exception-caught
358 logger.exception("Failed to save memory snapshot after training error.")
359 finally:
360 self.base.abort_profilers()
361 raise
363 synchronize()
364 self.base.destroy_distributed()