Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / llm_trainer.py: 0%
197 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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
18its ``_build_*`` methods selectively. Overrides ``_build_model_assets``,
19``_build_data_transform``, ``_build_dataset``, and ``_build_collate_fn``
20for complete real-data training pipeline.
22"""
23import logging
24from typing import Any, Dict, List
26import torch
27from torch.utils.data import Dataset
29from hyper_parallel.trainer.base import BaseTrainer
31logger = logging.getLogger(__name__)
34class LLMTrainer:
35 """Trainer for LM pretraining and SFT.
37 Composition pattern — calls BaseTrainer's _build_* methods in order,
38 overriding data pipeline steps for real tokenized data.
40 Supports:
41 - ``data.type = "dummy"``: random tokens for quick FSDP validation
42 - ``data.type = "hf_datasets"``: real HuggingFace datasets with tokenization
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() # 覆盖: 加载 tokenizer
56 self._build_data_transform() # 覆盖: tokenize 函数
57 self._build_dataset() # 覆盖: 真实数据集加载 + tokenize
58 self._build_collate_fn() # 覆盖: 支持变长 padding
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 For dummy data, tokenizer is not needed.
77 For real data, loads HF AutoTokenizer from ``model.weights_path``
78 or ``model.tokenizer_path``.
79 """
80 data_type = getattr(self.base.args.data, 'type', 'dummy')
81 if data_type == 'dummy':
82 self.base.tokenizer = None
83 return
85 # Try tokenizer_path first, fall back to weights_path
86 model_cfg = self.base.args.model
87 tokenizer_path = getattr(model_cfg, 'tokenizer_path', None)
88 if not tokenizer_path:
89 tokenizer_path = getattr(model_cfg, 'weights_path', None)
91 if not tokenizer_path:
92 raise ValueError(
93 "data.type='hf_datasets' requires model.tokenizer_path or "
94 "model.weights_path to load tokenizer."
95 )
97 from transformers import AutoTokenizer # pylint: disable=C0415 # optional dep
98 self.base.tokenizer = AutoTokenizer.from_pretrained(
99 tokenizer_path, trust_remote_code=True
100 )
101 # Ensure pad token exists
102 if self.base.tokenizer.pad_token is None:
103 self.base.tokenizer.pad_token = self.base.tokenizer.eos_token
104 logger.info("Tokenizer loaded: %s (vocab=%d)",
105 tokenizer_path, len(self.base.tokenizer))
107 def _build_data_transform(self):
108 """Build tokenization transform.
110 Creates a function that tokenizes raw text into input_ids + labels.
111 Labels are a copy of input_ids (causal LM: predict next token).
112 Prompt tokens can be masked with -100 for SFT.
113 """
114 if self.base.tokenizer is None:
115 self.base.data_transform = None
116 return
118 max_seq_len = getattr(self.base.args.data, 'max_seq_len', 2048)
119 tokenizer = self.base.tokenizer
120 text_key = getattr(self.base.args.data, 'text_key', 'text')
121 data_type = getattr(self.base.args.data, 'type', 'dummy')
122 template = getattr(self.base.args.data, 'template', 'empty')
124 def _tokenize_fn(examples):
125 """Tokenize text and create causal LM labels.
127 Supports:
128 - Plain text (text_key field)
129 - Alpaca format (instruction/input/output)
130 """
131 # SFT label masking: prompt tokens → IGNORE_INDEX, response
132 # tokens kept. Truncation prioritises the response side.
133 ignore_index = -100
135 def _infer_seqlen(s_len, t_len, cutoff):
136 if t_len * 2 < cutoff:
137 max_t = cutoff
138 elif s_len * 2 < cutoff:
139 max_t = cutoff - s_len
140 else:
141 max_t = int(cutoff * (t_len / (s_len + t_len)))
142 new_t = min(max_t, t_len)
143 max_s = max(cutoff - new_t, 0)
144 new_s = min(max_s, s_len)
145 return new_s, new_t
147 if "instruction" in examples and data_type == "json_file" and template == "empty":
148 instructions = examples["instruction"]
149 inputs = examples.get("input", [""] * len(instructions))
150 outputs = examples["output"]
151 result = {"input_ids": [], "labels": []}
152 for inst, inp, out in zip(instructions, inputs, outputs):
153 prompt_text = inst + (("\n" + inp) if inp else "")
154 prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"]
155 response_ids = tokenizer(out, add_special_tokens=False)["input_ids"]
156 s_len, t_len = _infer_seqlen(len(prompt_ids), len(response_ids), max_seq_len)
157 prompt_ids = prompt_ids[:s_len]
158 response_ids = response_ids[:t_len]
159 ids = prompt_ids + response_ids
160 labels = [ignore_index] * len(prompt_ids) + list(response_ids)
161 if len(ids) > 0:
162 result["input_ids"].append(ids)
163 result["labels"].append(labels)
164 return result
166 if "instruction" in examples and data_type == "json_file":
167 # Alpaca format with chat-style template (legacy default)
168 instructions = examples["instruction"]
169 inputs = examples.get("input", [""] * len(instructions))
170 outputs = examples["output"]
171 texts = []
172 for inst, inp, out in zip(instructions, inputs, outputs):
173 if inp:
174 texts.append(f"Human: {inst}\n{inp}\n\nAssistant: {out}")
175 else:
176 texts.append(f"Human: {inst}\n\nAssistant: {out}")
177 else:
178 # Plain text format
179 texts = examples[text_key]
180 if isinstance(texts, str):
181 texts = [texts]
183 tokenized = tokenizer(
184 texts,
185 truncation=True,
186 max_length=max_seq_len,
187 padding=False,
188 return_attention_mask=False,
189 )
191 result = {"input_ids": [], "labels": []}
192 for ids in tokenized["input_ids"]:
193 if len(ids) > 0:
194 result["input_ids"].append(ids)
195 result["labels"].append(ids.copy())
197 return result
199 self.base.data_transform = _tokenize_fn
200 logger.info("Data transform: tokenize max_seq_len=%d, format=%s",
201 max_seq_len, "alpaca" if data_type == "json_file" else text_key)
203 def _build_dataset(self):
204 """Build training dataset with full tokenization pipeline.
206 For dummy data, delegates to BaseTrainer.
207 For hf_datasets, loads + tokenizes + filters empty examples.
208 """
209 data_type = getattr(self.base.args.data, 'type', 'dummy')
211 if data_type == 'dummy':
212 self.base._build_dataset() # pylint: disable=protected-access
213 return
215 if data_type == 'preset_pt':
216 self._build_preset_pt_dataset()
217 return
219 if data_type not in ('hf_datasets', 'json_file'):
220 raise ValueError(
221 f"LLMTrainer supports data.type 'dummy', 'hf_datasets', 'json_file', or 'preset_pt', "
222 f"got '{data_type}'"
223 )
225 # pylint: disable=C0415
226 from datasets import load_dataset # pylint: disable=C0415 # optional dep
228 train_path = self.base.args.data.train_path
229 data_subset = getattr(self.base.args.data, 'subset', None)
231 logger.info("Loading dataset: type=%s, path=%s", data_type, train_path)
233 if data_type == 'json_file':
234 # Load local JSON file (alpaca format: instruction/input/output)
235 ds = load_dataset("json", data_files=train_path, split="train")
236 elif data_subset:
237 ds = load_dataset(train_path, data_subset, split="train")
238 else:
239 ds = load_dataset(train_path, split="train")
241 # Limit dataset size if specified
242 train_size = getattr(self.base.args.data, 'train_size', None)
243 if train_size and train_size < len(ds):
244 ds = ds.select(range(train_size))
245 logger.info("Dataset truncated to %d samples", train_size)
247 # Tokenize
248 if self.base.data_transform:
249 ds = ds.map(
250 self.base.data_transform,
251 batched=True,
252 remove_columns=ds.column_names,
253 desc="Tokenizing",
254 )
256 # Filter empty sequences
257 ds = ds.filter(lambda x: len(x["input_ids"]) > 0)
259 # Convert to torch tensors
260 class TokenizedDataset(torch.utils.data.Dataset):
261 """Wrap HF dataset for torch DataLoader."""
262 def __init__(self, hf_ds):
263 self.data = hf_ds
265 def __len__(self):
266 return len(self.data)
268 def __getitem__(self, idx):
269 item = self.data[idx]
270 return {
271 "input_ids": torch.tensor(item["input_ids"], dtype=torch.long),
272 "labels": torch.tensor(item["labels"], dtype=torch.long),
273 }
275 self.base.train_dataset = TokenizedDataset(ds)
276 self.base.state.max_steps = min(
277 self.base.args.train.max_steps,
278 len(self.base.train_dataset) // max(
279 self.base.args.train.global_batch_size, 1
280 ),
281 )
282 logger.info("Dataset ready: %d samples, max_steps=%d",
283 len(self.base.train_dataset), self.base.state.max_steps)
285 def _build_preset_pt_dataset(self):
286 """Load pre-tokenized batches from a .pt file (List[Dict[str, Tensor]]).
288 ``data.train_path`` is the .pt file. Each entry is a dict of tensors
289 with shape ``(global_batch, seq_len)``. The dataset returns a flat
290 sequence of per-sample dicts so the standard DataLoader can batch them.
291 Use this when the dataset has already been tokenized offline and the
292 token stream should be replayed deterministically.
293 """
294 # pylint: disable=C0415
296 train_path = self.base.args.data.train_path
297 if not train_path:
298 raise ValueError("data.train_path is required when data.type='preset_pt'")
299 batches = torch.load(train_path, map_location="cpu", weights_only=False)
300 if not isinstance(batches, list) or not batches:
301 raise ValueError(f"preset_pt expects List, got {type(batches)}")
302 per_sample = []
303 for b in batches:
304 # Two formats: stacked dict ``{input_ids: (B,S), labels: (B,S)}``,
305 # or list of per-rank dicts (preserves per-rank dynamic seq_len).
306 if isinstance(b, list):
307 for br in b:
308 ids = br["input_ids"]
309 labels = br["labels"]
310 attn = br.get("attention_mask")
311 for i in range(ids.shape[0]):
312 rec = {
313 "input_ids": ids[i].clone(),
314 "labels": labels[i].clone(),
315 }
316 if attn is not None and attn.dim() == 2:
317 rec["attention_mask"] = attn[i].clone()
318 per_sample.append(rec)
319 else:
320 ids = b["input_ids"]
321 labels = b["labels"]
322 attn = b.get("attention_mask")
323 for i in range(ids.shape[0]):
324 rec = {
325 "input_ids": ids[i].clone(),
326 "labels": labels[i].clone(),
327 }
328 if attn is not None and attn.dim() == 2:
329 rec["attention_mask"] = attn[i].clone()
330 per_sample.append(rec)
332 class PresetPtDataset(Dataset):
333 def __init__(self, samples):
334 self.samples = samples
336 def __len__(self):
337 return len(self.samples)
339 def __getitem__(self, idx):
340 return self.samples[idx]
342 self.base.train_dataset = PresetPtDataset(per_sample)
343 max_steps = getattr(self.base.args.train, "max_steps", None)
344 if max_steps:
345 self.base.state.max_steps = int(max_steps)
346 logger.info("preset_pt dataset: %d samples loaded from %s", len(per_sample), train_path)
348 def _build_collate_fn(self):
349 """Build collator with proper padding.
351 Pads input_ids with pad_token_id (or 0) and labels with -100.
352 """
354 pad_id = 0
355 if self.base.tokenizer and self.base.tokenizer.pad_token_id is not None:
356 pad_id = self.base.tokenizer.pad_token_id
358 def _lm_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
359 """Pad sequences to max length in batch."""
360 max_len = max(item["input_ids"].size(0) for item in batch)
361 input_ids_list = []
362 labels_list = []
364 for item in batch:
365 seq_len = item["input_ids"].size(0)
366 pad_len = max_len - seq_len
368 if pad_len > 0:
369 input_ids_list.append(
370 torch.cat([item["input_ids"],
371 torch.full((pad_len,), pad_id, dtype=torch.long)])
372 )
373 labels_list.append(
374 torch.cat([item["labels"],
375 torch.full((pad_len,), -100, dtype=torch.long)])
376 )
377 else:
378 input_ids_list.append(item["input_ids"])
379 labels_list.append(item["labels"])
381 return {
382 "input_ids": torch.stack(input_ids_list),
383 "labels": torch.stack(labels_list),
384 }
386 self.base.collate_fn = _lm_collate
388 # ------------------------------------------------------------------
389 # Delegated methods
390 # ------------------------------------------------------------------
392 def train(self):
393 """Delegate to BaseTrainer.train()."""
394 self.base.train()
396 def train_step(self, data_iterator):
397 """Delegate to BaseTrainer.train_step()."""
398 return self.base.train_step(data_iterator)