File size: 19,057 Bytes
93c8cdd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 | #!/usr/bin/env python3
"""Train the TenaOS single task-tagged Gemma 4 LoRA adapter with Unsloth."""
from __future__ import annotations
import argparse
import json
import math
import traceback
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
DEFAULT_BASE_MODEL = "unsloth/gemma-4-E4B-it"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--train-jsonl", type=Path, required=True)
parser.add_argument("--validation-jsonl", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--base-model", default=DEFAULT_BASE_MODEL)
parser.add_argument("--merge-output-dir", type=Path, help="Optional directory for merged weights.")
parser.add_argument(
"--max-seq-length",
type=int,
default=32768,
help=(
"Real multi-turn conversations run much longer than the old flattened single-completion "
"rows. Measured char-length distribution across the reconstructed corpus (rough proxy, "
"~3.5 chars/token): form median ~22.4K chars / p95 ~50.5K chars / max ~125.5K chars; "
"report median ~9.8K / p95 ~40.9K / max ~66K; cds/patient_education/scribe all comfortably "
"under 16K chars. 32768 tokens covers the large majority without the extreme few-outlier "
"tail; those are truncated by SFTConfig's max_length rather than raising the default so high "
"that most batches waste compute on padding."
),
)
parser.add_argument("--epochs", type=float, default=3.0)
parser.add_argument("--max-steps", type=int, default=-1, help="Optional cap for smoke tests or short runs.")
parser.add_argument("--per-device-train-batch-size", type=int, default=1)
parser.add_argument("--gradient-accumulation-steps", type=int, default=8)
parser.add_argument("--learning-rate", type=float, default=2e-4)
parser.add_argument("--warmup-ratio", type=float, default=0.03)
parser.add_argument("--logging-steps", type=int, default=10)
parser.add_argument("--eval-steps", type=int, default=100)
parser.add_argument("--save-steps", type=int, default=250)
parser.add_argument("--seed", type=int, default=3407)
parser.add_argument("--load-in-4bit", action="store_true", help="Use QLoRA fallback if BF16 LoRA does not fit.")
parser.add_argument("--dataset-num-proc", type=int, default=4)
parser.add_argument("--packing", action="store_true", help="Enable TRL packing for short examples.")
parser.add_argument(
"--chat-template",
default=None,
help=(
"Optional Unsloth canned chat_template name (see unsloth.chat_templates.CHAT_TEMPLATES). "
"Defaults to the base model's own tokenizer chat_template (recommended: this is the same "
"template baked into the merged GGUF that llama.cpp renders with --jinja at serve time, "
"so training and inference stay byte-identical without depending on Unsloth shipping a "
"matching canned template for this specific base model)."
),
)
parser.add_argument(
"--instruction-part",
default="<|turn>user\n",
help=(
"Turn-start marker for non-assistant turns, used by train_on_responses_only masking. "
"NOTE: this Gemma 4 base model's chat template uses '<|turn>role ... <turn|>' markers, "
"NOT the classic Gemma 2/3 '<start_of_turn>role ... <end_of_turn>' tokens — verified by "
"rendering a sample conversation with the tokenizer's own apply_chat_template. Using the "
"wrong markers silently zeroes out every example's loss mask (train_on_responses_only "
"finds no matches), which crashes the RandomSampler with num_samples=0."
),
)
parser.add_argument(
"--response-part",
default="<|turn>model\n",
help="Turn-start marker for assistant turns, used by train_on_responses_only masking.",
)
parser.add_argument(
"--no-response-only-masking",
action="store_true",
help="Disable response-only loss masking (NOT recommended — trains loss over tool/user tokens too).",
)
parser.add_argument(
"--use-flex-attention",
action="store_true",
help=(
"Force torch's flex_attention backend instead of Unsloth's default sdpa. NOT "
"recommended: under Unsloth's custom gradient checkpointing the layer recompute runs "
"outside torch.compile, so flex_attention's backward dispatches to its eager dense "
"decomposition (torch/_higher_order_ops/flex_attention.py::sdpa_dense_backward), which "
"materializes full quadratic (B, H, N, N) score/grad tensors — ~9 GiB each at "
"N=24576 with 8 heads — and reliably OOMs an 80 GB A100 on the first backward pass. "
"The original rationale (head_dim 512 exceeding fused-kernel limits) was wrong: this "
"model's config has head_dim=256, which sdpa's fused flash/mem-efficient kernels fully "
"support."
),
)
parser.add_argument(
"--no-length-filter",
action="store_true",
help=(
"Disable dropping examples whose rendered token length exceeds --max-seq-length. By default "
"such outliers are dropped entirely rather than silently truncated by SFTConfig's max_length, "
"since truncation can cut a conversation off mid tool-call/JSON payload and teach the model a "
"corrupted, incomplete turn. Measured on the current corpus, only ~0.1% of examples (the "
"extreme long tail of the 'form' task) exceed 24576 tokens, so dropping them is cheap."
),
)
return parser.parse_args()
def main() -> None:
args = parse_args()
args.output_dir.mkdir(parents=True, exist_ok=True)
try:
from unsloth import FastVisionModel
from unsloth.chat_templates import CHAT_TEMPLATES, get_chat_template, train_on_responses_only
import torch
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
except ImportError as exc:
raise SystemExit(
"Missing training dependencies. Install Unsloth, TRL, datasets, transformers, peft, and torch "
"in the training environment before running this script."
) from exc
model, tokenizer = FastVisionModel.from_pretrained(
args.base_model,
load_in_4bit=args.load_in_4bit,
use_gradient_checkpointing="unsloth",
max_seq_length=args.max_seq_length,
)
# Keep Unsloth's default sdpa, but force the repeat-KV path instead of enable_gqa.
#
# Verified by intercepting the actual SDPA calls during a 24576-token forward:
# - sliding-window layers (35/42): q/k (1, 8, N, 256) + boolean mask
# -> fused memory-efficient kernel, linear memory. Fine as-is.
# - global full-attention layers (7/42): q (1, 8, N, 512), k/v (1, 2, N, 512)
# with enable_gqa=True and is_causal.
# No fused SDPA kernel supports head_dim=512 *with* enable_gqa (flash caps at
# head_dim 256; the memory-efficient kernel rejects enable_gqa), so those 7 layers
# silently dispatch to the quadratic-memory math kernel, whose backward materializes
# an (8, N, N) fp32 tensor -- an 18 GiB single allocation at N=24576, which OOMs the
# 80 GB A100. With K/V repeated to the full 8 heads the memory-efficient kernel
# accepts head_dim=512 (measured: 2.49 GiB peak for fwd+bwd of one such layer at
# N=24576, vs OOM before). repeat_kv is mathematically identical to enable_gqa,
# so this changes memory/speed only, never outputs.
#
# (A previous revision instead forced flex_attention here, which OOMed even harder:
# under Unsloth's gradient checkpointing the recompute runs uncompiled, and eager
# flex_attention's backward decomposes to dense sdpa_dense_backward with full
# (B, H, N, N) score tensors on *every* layer.)
if args.use_flex_attention:
model.set_attn_implementation("flex_attention")
print(f"Attention implementation forced to: {model.config._attn_implementation}")
else:
import transformers.integrations.sdpa_attention as _sdpa_integration
_sdpa_integration.use_gqa_in_sdpa = lambda attention_mask, key: False
print(
"SDPA enable_gqa disabled (forcing repeat_kv) so the 512-dim global-attention "
"layers use the fused memory-efficient kernel instead of the quadratic math fallback."
)
if args.chat_template:
if args.chat_template not in CHAT_TEMPLATES:
raise SystemExit(
f"--chat-template={args.chat_template!r} is not a known Unsloth canned template "
f"(available: {sorted(CHAT_TEMPLATES.keys())}). Omit --chat-template to use the base "
"model's own tokenizer chat_template instead (recommended)."
)
tokenizer = get_chat_template(tokenizer, chat_template=args.chat_template)
elif not getattr(tokenizer, "chat_template", None):
raise SystemExit(
f"Base model {args.base_model!r} has no tokenizer chat_template and no --chat-template "
"override was given. A previous run hardcoded chat_template='gemma-4-thinking', which does "
"not exist in this Unsloth version's CHAT_TEMPLATES registry and would crash here — pass an "
"explicit --chat-template only if you have verified it renders tool-role turns the same way "
"llama.cpp's --jinja rendering of the merged GGUF does (see the template-check step)."
)
print(f"Using chat_template (first 200 chars): {str(tokenizer.chat_template)[:200]!r}")
model = FastVisionModel.get_peft_model(
model,
finetune_vision_layers=False,
finetune_language_layers=True,
finetune_attention_modules=True,
finetune_mlp_modules=True,
r=16,
lora_alpha=32,
# NOTE: the installed peft version's MoE ParamWrapper path (used for
# `experts.gate_up_proj` / `experts.down_proj`) raises
# `ValueError: lora.ParamWrapper does not work with lora_dropout != 0`
# for this model's grouped-expert parameters, so dropout must stay 0.
lora_dropout=0.0,
bias="none",
random_state=args.seed,
use_rslora=False,
loftq_config=None,
target_modules="all-linear",
)
dataset = load_dataset(
"json",
data_files={
"train": str(args.train_jsonl),
"validation": str(args.validation_jsonl),
},
)
def formatting_prompts_func(examples: dict[str, Any]) -> dict[str, list[str]]:
texts = [
tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False).removeprefix("<bos>")
for convo in examples["conversations"]
]
return {"text": texts}
dataset = dataset.map(
formatting_prompts_func,
batched=True,
num_proc=args.dataset_num_proc,
remove_columns=dataset["train"].column_names,
)
if not args.no_length_filter:
def _add_token_length(examples: dict[str, Any]) -> dict[str, list[int]]:
# Gemma4's tokenizer is a full multimodal processor; it requires `text=`
# passed as an explicit keyword (positional args get swallowed as `images`
# by Unsloth's patched __call__ wrapper, leaving text=None internally).
# padding=False is required too: batched calls otherwise pad every row in
# the (default size-1000) map batch out to that batch's single longest
# sequence, which silently inflates the reported length of every short
# example that happens to share a batch with one long outlier.
return {
"_n_tokens": [
len(ids)
for ids in tokenizer(
text=examples["text"], add_special_tokens=False, padding=False
)["input_ids"]
]
}
for split in ("train", "validation"):
before = len(dataset[split])
dataset[split] = dataset[split].map(_add_token_length, batched=True, num_proc=args.dataset_num_proc)
dataset[split] = dataset[split].filter(
lambda ex: ex["_n_tokens"] <= args.max_seq_length,
num_proc=args.dataset_num_proc,
)
dataset[split] = dataset[split].remove_columns(["_n_tokens"])
after = len(dataset[split])
print(
f"Length filter ({split}): dropped {before - after}/{before} examples exceeding "
f"--max-seq-length={args.max_seq_length} tokens (kept {after})."
)
else:
print(
"WARNING: --no-length-filter set; over-length examples will be silently truncated by "
"SFTConfig's max_length instead of being dropped."
)
train_size = len(dataset["train"])
effective_batch = max(1, args.per_device_train_batch_size * args.gradient_accumulation_steps)
estimated_steps = max(1, math.ceil(train_size / effective_batch) * max(1, math.ceil(args.epochs)))
warmup_steps = max(1, int(estimated_steps * args.warmup_ratio))
train_args = SFTConfig(
output_dir=str(args.output_dir),
num_train_epochs=args.epochs,
max_steps=args.max_steps,
per_device_train_batch_size=args.per_device_train_batch_size,
gradient_accumulation_steps=args.gradient_accumulation_steps,
learning_rate=args.learning_rate,
warmup_steps=warmup_steps,
logging_steps=args.logging_steps,
# In-loop evaluation is disabled: Unsloth's patched unsloth_prediction_step
# (unsloth/models/rl.py) unconditionally sets UNSLOTH_RETURN_LOGITS=1 and
# calls compute_loss(..., return_outputs=True) during eval, ignoring
# prediction_loss_only entirely. That materializes the full
# (batch, seq_len, ~262K vocab) logits tensor and casts it to fp32 inside
# the forward pass itself -- verified to OOM (41.6 GiB single allocation)
# on this model's vocab size regardless of TrainingArguments settings.
# prediction_loss_only only affects what the Trainer keeps afterward, by
# which point the OOM has already happened. Real validation/quality
# checks happen post-training via evaluate_lora_ab.py against the held-out
# test split instead; save_steps below still gives periodic checkpoints.
eval_strategy="no",
save_steps=args.save_steps,
save_total_limit=3,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=args.seed,
bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
fp16=torch.cuda.is_available() and not torch.cuda.is_bf16_supported(),
dataset_num_proc=args.dataset_num_proc,
dataset_text_field="text",
max_length=args.max_seq_length,
packing=args.packing,
report_to="none",
)
trainer = SFTTrainer(
model=model,
processing_class=tokenizer,
train_dataset=dataset["train"],
eval_dataset=dataset["validation"],
args=train_args,
)
if not args.no_response_only_masking:
# Every training example is now a real multi-turn conversation
# (system/user/assistant tool-calls/tool results) reconstructed from
# production traces, not a single flattened completion. Without this,
# SFT loss would run over system/user/tool tokens too — including
# tool results the model never generates itself — which both wastes
# signal and actively trains the model to try to "predict" search
# results and prior context verbatim. Mask to assistant-authored
# tokens only, matched via the chat template's own turn markers.
trainer = train_on_responses_only(
trainer,
instruction_part=args.instruction_part,
response_part=args.response_part,
)
print(
f"Response-only masking enabled (instruction_part={args.instruction_part!r}, "
f"response_part={args.response_part!r})."
)
else:
print("WARNING: response-only masking disabled via --no-response-only-masking.")
train_result = trainer.train()
trainer.save_model(str(args.output_dir))
tokenizer.save_pretrained(str(args.output_dir))
metadata = {
"schema_version": "tenaos_lora_training_run_v1",
"created_at": datetime.now(timezone.utc).isoformat(),
"base_model": args.base_model,
"train_jsonl": str(args.train_jsonl),
"validation_jsonl": str(args.validation_jsonl),
"output_dir": str(args.output_dir),
"load_in_4bit": args.load_in_4bit,
"chat_template_name": args.chat_template or "<tokenizer_native>",
"chat_template_prefix": str(tokenizer.chat_template)[:200],
"response_only_masking": not args.no_response_only_masking,
"instruction_part": args.instruction_part,
"response_part": args.response_part,
"hyperparameters": {
"lora_rank": 16,
"lora_alpha": 32,
"lora_dropout": 0.0,
"epochs": args.epochs,
"max_steps": args.max_steps,
"max_seq_length": args.max_seq_length,
"per_device_train_batch_size": args.per_device_train_batch_size,
"gradient_accumulation_steps": args.gradient_accumulation_steps,
"learning_rate": args.learning_rate,
"warmup_steps": warmup_steps,
"seed": args.seed,
"target_modules": "all-linear",
"finetune_language_layers": True,
"finetune_attention_modules": True,
"finetune_mlp_modules": True,
"finetune_vision_layers": False,
},
"train_metrics": train_result.metrics,
}
(args.output_dir / "training_metadata.json").write_text(
json.dumps(metadata, indent=2, ensure_ascii=False, sort_keys=True) + "\n",
encoding="utf-8",
)
if args.merge_output_dir:
args.merge_output_dir.mkdir(parents=True, exist_ok=True)
model.save_pretrained_merged(str(args.merge_output_dir), tokenizer, save_method="merged_16bit")
metadata["merge_output_dir"] = str(args.merge_output_dir)
(args.output_dir / "training_metadata.json").write_text(
json.dumps(metadata, indent=2, ensure_ascii=False, sort_keys=True) + "\n",
encoding="utf-8",
)
print(json.dumps(metadata, indent=2, ensure_ascii=False, sort_keys=True))
if __name__ == "__main__":
try:
main()
except BaseException:
traceback.print_exc()
raise
|