Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / llm_trainer.py: 0%
141 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""LLMTrainer — Language Model pretraining and SFT.
17Holds a ``BaseTrainer`` instance and calls its ``_build_*`` methods
18selectively. Overrides ``_build_model_assets``, ``_build_data_transform``
19and ``_build_collate_fn``; dataset construction is delegated to the
20shared :func:`hyper_parallel.data.build_dataset` registry.
21"""
22import logging
23from typing import Any, Dict, List
25import torch
27from hyper_parallel.trainer.base import BaseTrainer
29logger = logging.getLogger(__name__)
32class LLMTrainer:
33 """Trainer for LM pretraining and SFT.
35 Composition pattern — calls BaseTrainer's _build_* methods in order,
36 overriding data pipeline steps for real tokenized data.
38 Supports every ``data.type`` registered with
39 :data:`hyper_parallel.data.DATASET_REGISTRY` — built-in formats are
40 ``dummy`` (random tokens), ``hf_datasets`` / ``json_file`` (HF +
41 Alpaca), ``preset_pt`` (replayed batches), and ``megatron`` (Megatron
42 ``.bin``/``.idx``).
44 Args:
45 args: Training configuration parsed from YAML.
46 """
48 def __init__(self, args):
49 self.base = BaseTrainer(args)
51 # 13 steps — call base's methods, override where needed
52 self.base._setup()
53 self.base._build_model()
54 self.base._freeze_model()
55 self._build_model_assets()
56 self._build_data_transform()
57 self.base._build_dataset()
58 self._build_collate_fn()
59 self.base._build_dataloader()
60 self.base._build_parallelized_model()
61 self.base._build_optimizer()
62 self.base._build_lr_scheduler()
63 self.base._build_training_context()
64 self.base._init_callbacks()
65 # Fire one-shot ``on_init_end`` AFTER every ``_build_*`` — this is
66 # the canonical "trainer is fully built" lifecycle hook.
67 self.base.on_init_end()
69 # ------------------------------------------------------------------
70 # Overridden _build_* methods
71 # ------------------------------------------------------------------
73 def _build_model_assets(self):
74 """Build tokenizer for data processing.
76 Pre-tokenized formats (``dummy`` random tokens, ``megatron`` .bin/.idx,
77 ``preset_pt`` replayed tensors) carry no raw text, so no tokenizer is
78 built for them. Text formats (``hf_datasets`` / ``json_file``) load an
79 HF AutoTokenizer from ``model.tokenizer_path`` or ``model.weights_path``.
80 """
81 data_type = self.base.args.data.type
82 if data_type in ('dummy', 'megatron', 'preset_pt'):
83 self.base.tokenizer = None
84 return
86 # Try tokenizer_path first, fall back to weights_path
87 model_cfg = self.base.args.model
88 tokenizer_path = model_cfg.tokenizer_path
89 if not tokenizer_path:
90 tokenizer_path = model_cfg.weights_path
92 if not tokenizer_path:
93 raise ValueError(
94 "data.type='hf_datasets' requires model.tokenizer_path or "
95 "model.weights_path to load tokenizer."
96 )
98 from transformers import AutoTokenizer # pylint: disable=C0415 # optional dep
99 self.base.tokenizer = AutoTokenizer.from_pretrained(
100 tokenizer_path, trust_remote_code=True
101 )
102 # Ensure pad token exists
103 if self.base.tokenizer.pad_token is None:
104 self.base.tokenizer.pad_token = self.base.tokenizer.eos_token
105 logger.info("Tokenizer loaded: %s (vocab=%d)",
106 tokenizer_path, len(self.base.tokenizer))
108 def _build_data_transform(self):
109 """Build tokenization transform.
111 Creates a function that tokenizes raw text into input_ids + labels.
112 Labels are a copy of input_ids (causal LM: predict next token).
113 Prompt tokens can be masked with -100 for SFT.
114 """
115 if self.base.tokenizer is None:
116 self.base.data_transform = None
117 return
119 max_seq_len = self.base.args.data.max_seq_len
120 tokenizer = self.base.tokenizer
121 text_key = self.base.args.data.text_key
122 data_type = self.base.args.data.type
123 template = self.base.args.data.template
125 def _tokenize_fn(examples):
126 """Tokenize text and create causal LM labels.
128 Supports:
129 - Plain text (text_key field)
130 - Alpaca format (instruction/input/output)
131 """
132 # SFT label masking: prompt tokens → IGNORE_INDEX, response
133 # tokens kept. Truncation prioritises the response side.
134 ignore_index = -100
136 def _infer_seqlen(s_len, t_len, cutoff):
137 if t_len * 2 < cutoff:
138 max_t = cutoff
139 elif s_len * 2 < cutoff:
140 max_t = cutoff - s_len
141 else:
142 max_t = int(cutoff * (t_len / (s_len + t_len)))
143 new_t = min(max_t, t_len)
144 max_s = max(cutoff - new_t, 0)
145 new_s = min(max_s, s_len)
146 return new_s, new_t
148 if "instruction" in examples and data_type == "json_file" and template == "empty":
149 instructions = examples["instruction"]
150 inputs = examples.get("input", [""] * len(instructions))
151 outputs = examples["output"]
152 result = {"input_ids": [], "labels": []}
153 for inst, inp, out in zip(instructions, inputs, outputs):
154 prompt_text = inst + (("\n" + inp) if inp else "")
155 prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"]
156 response_ids = tokenizer(out, add_special_tokens=False)["input_ids"]
157 s_len, t_len = _infer_seqlen(len(prompt_ids), len(response_ids), max_seq_len)
158 prompt_ids = prompt_ids[:s_len]
159 response_ids = response_ids[:t_len]
160 ids = prompt_ids + response_ids
161 labels = [ignore_index] * len(prompt_ids) + list(response_ids)
162 if len(ids) > 0:
163 result["input_ids"].append(ids)
164 result["labels"].append(labels)
165 return result
167 if "instruction" in examples and data_type == "json_file":
168 # Alpaca format with chat-style template (legacy default)
169 instructions = examples["instruction"]
170 inputs = examples.get("input", [""] * len(instructions))
171 outputs = examples["output"]
172 texts = []
173 for inst, inp, out in zip(instructions, inputs, outputs):
174 if inp:
175 texts.append(f"Human: {inst}\n{inp}\n\nAssistant: {out}")
176 else:
177 texts.append(f"Human: {inst}\n\nAssistant: {out}")
178 else:
179 # Plain text format
180 texts = examples[text_key]
181 if isinstance(texts, str):
182 texts = [texts]
184 tokenized = tokenizer(
185 texts,
186 truncation=True,
187 max_length=max_seq_len,
188 padding=False,
189 return_attention_mask=False,
190 )
192 result = {"input_ids": [], "labels": []}
193 for ids in tokenized["input_ids"]:
194 if len(ids) > 0:
195 result["input_ids"].append(ids)
196 result["labels"].append(ids.copy())
198 return result
200 self.base.data_transform = _tokenize_fn
201 logger.info("Data transform: tokenize max_seq_len=%d, format=%s",
202 max_seq_len, "alpaca" if data_type == "json_file" else text_key)
204 def _build_collate_fn(self):
205 """Build collator with proper padding.
207 Pads input_ids with pad_token_id (or 0) and labels with -100.
208 SequenceParallel TP and context parallel slice the sequence dim, so
209 variable-length batches additionally pad up to a multiple of
210 ``cp * tp``; the pad rides label ``-100`` and is masked out of the CE.
211 """
213 pad_id = 0
214 if self.base.tokenizer and self.base.tokenizer.pad_token_id is not None:
215 pad_id = self.base.tokenizer.pad_token_id
216 seq_divisor = self.base.parallel_dims.seq_divisor
218 def _lm_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
219 """Pad sequences to max length in batch."""
220 max_len = max(item["input_ids"].size(0) for item in batch)
221 if seq_divisor > 1 and max_len % seq_divisor:
222 max_len += seq_divisor - max_len % seq_divisor
223 input_ids_list = []
224 labels_list = []
226 for item in batch:
227 seq_len = item["input_ids"].size(0)
228 pad_len = max_len - seq_len
230 if pad_len > 0:
231 input_ids_list.append(
232 torch.cat([item["input_ids"],
233 torch.full((pad_len,), pad_id, dtype=torch.long)])
234 )
235 labels_list.append(
236 torch.cat([item["labels"],
237 torch.full((pad_len,), -100, dtype=torch.long)])
238 )
239 else:
240 input_ids_list.append(item["input_ids"])
241 labels_list.append(item["labels"])
243 out = {
244 "input_ids": torch.stack(input_ids_list),
245 "labels": torch.stack(labels_list),
246 }
247 if "num_items_in_batch" in batch[0]:
248 out["num_items_in_batch"] = sum(
249 int(item["num_items_in_batch"]) for item in batch
250 )
251 if "attention_mask" in batch[0]:
252 masks = []
253 for item in batch:
254 pad_len = max_len - item["attention_mask"].size(0)
255 masks.append(torch.nn.functional.pad(item["attention_mask"], (0, pad_len), value=0))
256 out["attention_mask"] = torch.stack(masks)
257 if "position_ids" in batch[0]:
258 positions = []
259 for item in batch:
260 pos = item["position_ids"]
261 pad_len = max_len - pos.shape[-1]
262 positions.append(torch.nn.functional.pad(pos, (0, pad_len), value=0))
263 if positions[0].dim() == 1:
264 out["position_ids"] = torch.stack(positions)
265 else:
266 out["position_ids"] = torch.stack(positions).transpose(0, 1).contiguous()
267 return out
269 self.base.collate_fn = _lm_collate
271 # ------------------------------------------------------------------
272 # Delegated methods
273 # ------------------------------------------------------------------
275 def train(self):
276 """Delegate to BaseTrainer.train()."""
277 self.base.train()
279 def train_step(self, data_iterator):
280 """Delegate to BaseTrainer.train_step()."""
281 return self.base.train_step(data_iterator)