#!/usr/bin/env python3 """ Phase 3 — QLoRA fine-tune of Muse-Glimmer-30B (text-only, perception encoder frozen). Run entirely inside WSL, from the venv at ~/glimmer/venv. See scripts/TRAINING.md for exact launch / pause / resume / monitor commands. Requirements this implements (see briefs/task-3-brief.md for the authoritative spec): 1. Load via Unsloth (FastLanguageModel, falling back to FastVisionModel) from the local HF cache only -- never re-downloads, asserts a cache hit up front. 2. Freezes the perception encoder (vision tower + vision->text projector): prefers Unsloth's own finetune_vision_layers=False-style flag, then independently verifies via named_parameters() that zero trainable params carry a vision/ projector-ish name. 3. LoRA r=16/alpha=16/dropout=0 on the text tower's attention + MLP projections only (q/k/v/o, gate/up/down) -- no embeddings, no lm_head. 4. Formats dataset/dataset.jsonl with tokenizer.apply_chat_template (bespoke <|start|>/<|message|>/<|eot|> template, reasoning_strength='high'), never hand-rolled. 5. Masks non-assistant turns via unsloth_zoo's train_on_responses_only, using the template's real turn markers (derived from chat_template.jinja + confirmed against tokenizer_config.json's response_template, and verified for real against 2 decoded dataset samples -- see briefs/task-3-report.md). 6. Trainer config is plan-mandated (see brief); checkpointing is tightened to save_steps=200 / save_total_limit=3 for pausability. 7-9. Pausable training: a TrainerCallback polls ~/glimmer/PAUSE on_step_end, saves a checkpoint, logs "PAUSED at step N", stops cleanly, and deletes the sentinel itself. --resume picks up the latest checkpoint (weights + optimizer + scheduler + step). 10. GPU strategy / OOM ladder is a launch-time concern (CUDA_VISIBLE_DEVICES / batch size / device_map), documented in TRAINING.md -- this script exposes the knobs as CLI flags rather than hardcoding one config. 11-13. Smoke leg: actually run at --max_steps 20 --max_seq_length 2048 (rung (c) -- rung (a)'s seq 4096 measured ~2550s/step and was abandoned as impractical; 100 steps was the original target but a long resumed run was killed by something external to this script partway through -- see the report for the full, honest chain of real-world adaptations). Loss series, one real pause/resume cycle with verified step continuity, 3 held-out generations, and a timing projection are all in briefs/task-3-report.md with real numbers. RESOLVED BLOCKER (was open when this script was first written -- see briefs/task-3-report.md for full history): `transformers==5.5.0` did not register the `muse_glimmer` model_type in CONFIG_MAPPING_NAMES at all. The orchestrator upgraded the shared venv to `transformers==5.15.0`, which does register it; `FastLanguageModel. from_pretrained` now loads the model (confirmed live, 22.2GB weights-only on one 3090). One live gotcha this upgrade surfaced: the tokenizer object this model's `from_pretrained` returns is a MuseGlimmerProcessor, not a plain tokenizer -- see `get_text_tokenizer()`. """ from __future__ import annotations import argparse import json import os import sys import time from pathlib import Path # --------------------------------------------------------------------------- # Environment -- must happen before any HF/transformers/unsloth import touches # the network. HF_HUB_OFFLINE=1 makes any cache-miss a loud, immediate error # instead of a silent download, which is how "assert the cache hit" is enforced. # --------------------------------------------------------------------------- GLIMMER_HOME = Path(os.environ.get("GLIMMER_HOME", os.path.expanduser("~/glimmer"))) os.environ.setdefault("HF_HOME", str(GLIMMER_HOME / "hf_home")) os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TOKENIZERS_PARALLELISM"] = "false" MODEL_NAME = "meta-models/Muse-Glimmer-30B" DATASET_PATH_DEFAULT = GLIMMER_HOME / "dataset" / "dataset.jsonl" OUTPUT_DIR_DEFAULT = GLIMMER_HOME / "runs" / "sentry-v01" PAUSE_SENTINEL_DEFAULT = GLIMMER_HOME / "PAUSE" HOLDOUT_DIR_DEFAULT = Path("/mnt/c/Users/Dwain-Admin/Desktop/GLIMMER-SENTRY-30B/dataset/holdout") # Bespoke chat template's exact assistant-turn markers. Derived from the model # repo's own chat_template.jinja (the `elif role == 'assistant'` branch, plain # reply case: recipient defaults to 'user', so a normal assistant turn renders # as `<|start|>assistant to=user<|message|>{content}<|eot|>`), and cross-checked # against tokenizer_config.json's structured `response_template` field, whose # `start_anchor` ('<|start|>assistant') + `fields.content.open_pattern` # ('to=user<\|message\|>') combine to the exact same string. Verified for real # (not just derived) against 2 decoded dataset examples -- see the report. INSTRUCTION_PART = "<|start|>user<|message|>" RESPONSE_PART = "<|start|>assistant to=user<|message|>" REASONING_STRENGTH = "high" # Name fragments that identify vision-tower / projector parameters regardless of # the exact attribute path Muse-Glimmer's implementation uses (config.json gives # `vision_config`/`muse_glimmer_vision`, ~1.8B params/50 layers/hidden 1536, and a # `projector_hidden_size`/`projector_hidden_act` pair for the vision->text # projector, but not the Python attribute name -- so this matches broadly and the # script prints every matched prefix it actually found for a human to sanity-check). VISION_NAME_FRAGMENTS = [ "vision", "visual", "projector", "vision_tower", "multi_modal_projector", "image_newline", "patch_embed", "vit.", ".vit", "perceiver", ] LORA_R = 16 LORA_ALPHA = 16 LORA_DROPOUT = 0.0 # Explicit fallback target-module leaf names if we must call get_peft_model without # Unsloth's vision-aware flags (e.g. if FastLanguageModel, not FastVisionModel, is # what actually loads this checkpoint). Attention + MLP projections only; no # embeddings, no lm_head. TEXT_TOWER_TARGET_MODULES = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ] def log(msg: str) -> None: print(f"[train.py {time.strftime('%H:%M:%S')}] {msg}", flush=True) def get_text_tokenizer(tokenizer_or_processor): """FastLanguageModel.from_pretrained returns a MuseGlimmerProcessor for this model (confirmed live, transformers 5.15.0), not a plain tokenizer -- its __call__'s first positional arg is `images`, so calling it positionally with text (`tokenizer(text)`) raises a base64/"Incorrect padding" error deep in the image-processing path instead of tokenizing. The plain tokenizer used for every text-only operation here (apply_chat_template, train_on_responses_only, the trainer's processing_class, generation) is the inner `.tokenizer` attribute. Falls back to the object itself if it's already a plain tokenizer (no `.tokenizer` attribute) so this is safe either way.""" return getattr(tokenizer_or_processor, "tokenizer", tokenizer_or_processor) # --------------------------------------------------------------------------- # 0. Cache-hit assertion # --------------------------------------------------------------------------- def assert_cache_hit(model_name: str) -> str: """Resolve the model snapshot from the local HF cache only. Raises loudly (via HF_HUB_OFFLINE=1, already set above) instead of downloading anything.""" from huggingface_hub import snapshot_download try: path = snapshot_download(model_name, local_files_only=True) except Exception as e: raise RuntimeError( f"Cache-hit assertion FAILED for '{model_name}' under HF_HOME=" f"{os.environ['HF_HOME']}. Refusing to download (HF_HUB_OFFLINE=1). " f"Original error: {e}" ) from e log(f"Cache hit confirmed: {model_name} -> {path}") return path # --------------------------------------------------------------------------- # 1-2. Model + tokenizer loading, perception-encoder freeze # --------------------------------------------------------------------------- def load_model_and_tokenizer(max_seq_length: int, device_map: str): """Try FastLanguageModel first, fall back to FastVisionModel, per the brief. Returns (model, tokenizer, loader_name). Raises with the exact upstream error (never falls back to raw transformers+peft) if both fail -- that is a BLOCKED condition per the brief, not something to improvise around.""" import torch from unsloth import FastLanguageModel, FastVisionModel common_kwargs = dict( model_name=MODEL_NAME, max_seq_length=max_seq_length, load_in_4bit=True, dtype=torch.bfloat16, device_map=device_map, ) log("Attempting FastLanguageModel.from_pretrained ...") try: model, tokenizer = FastLanguageModel.from_pretrained(**common_kwargs) return model, tokenizer, "FastLanguageModel" except Exception as e_lang: log(f"FastLanguageModel failed: {e_lang!r}") log("Falling back to FastVisionModel.from_pretrained ...") try: model, tokenizer = FastVisionModel.from_pretrained(**common_kwargs) return model, tokenizer, "FastVisionModel" except Exception as e_vision: raise RuntimeError( "BLOCKED: Muse-Glimmer-30B could not be loaded by either " "FastLanguageModel or FastVisionModel in unsloth 2026.8.12 / " f"transformers {__import__('transformers').__version__}.\n\n" f"FastLanguageModel error:\n{e_lang!r}\n\n" f"FastVisionModel error:\n{e_vision!r}\n\n" "Per the brief: do not improvise with raw transformers+peft here -- " "this is a BLOCKED condition to report, not to route around." ) from e_vision def apply_lora_with_frozen_vision(model, loader_name: str, use_gradient_checkpointing="unsloth"): """Attach LoRA to the text tower only. Prefers Unsloth's own finetune_vision_layers=False-style flag (FastVisionModel.get_peft_model) -- tried regardless of which loader actually succeeded, since that flag operates on the already-loaded model object via get_peft_regex's own module introspection, not on loader-stamped state. Falls back to the explicit text-tower target_modules list (FastLanguageModel.get_peft_model / plain LoraConfig) if the flag-based call raises (e.g. genuine incompatibility with a FastLanguageModel-loaded object). named_parameters() is independently verified afterward regardless of which path was used -- see verify_freeze().""" from unsloth import FastLanguageModel, FastVisionModel try: log("Attempting FastVisionModel.get_peft_model with finetune_vision_layers=False " "(Unsloth's native vision-freeze flag) ...") model = FastVisionModel.get_peft_model( model, r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, bias="none", finetune_vision_layers=False, finetune_language_layers=True, finetune_attention_modules=True, finetune_mlp_modules=True, use_gradient_checkpointing=use_gradient_checkpointing, random_state=42, ) log("FastVisionModel.get_peft_model succeeded.") except Exception as e: log(f"FastVisionModel.get_peft_model failed ({e!r}); falling back to the " f"explicit text-tower target_modules list via FastLanguageModel.get_peft_model.") model = FastLanguageModel.get_peft_model( model, r=LORA_R, target_modules=TEXT_TOWER_TARGET_MODULES, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, bias="none", use_gradient_checkpointing=use_gradient_checkpointing, random_state=42, ) return model def verify_freeze(model) -> dict: """Independent, from-scratch verification (not trusting the flag above): walk named_parameters(), tally total/trainable, and confirm zero trainable params carry a vision/projector-ish name. Defensively sets requires_grad=False on any that slip through, and prints + returns everything for the report.""" total_params = 0 trainable_params = 0 trainable_vision_params = 0 matched_prefixes = set() offending_names = [] for name, p in model.named_parameters(): n = p.numel() total_params += n lname = name.lower() is_vision_ish = any(frag in lname for frag in VISION_NAME_FRAGMENTS) if p.requires_grad: trainable_params += n if is_vision_ish: trainable_vision_params += n offending_names.append(name) matched_prefixes.add(".".join(name.split(".")[:4])) # Defensive: the brief asks to freeze regardless of what the flag did. p.requires_grad_(False) trainable_params -= n result = { "total_params": total_params, "trainable_params": trainable_params, "trainable_vision_params_before_defensive_freeze": trainable_vision_params, "offending_names_sample": offending_names[:20], } log(f"Freeze verification: total_params={total_params:,} " f"trainable_params={trainable_params:,} " f"trainable_pct={100*trainable_params/max(total_params,1):.4f}%") if trainable_vision_params > 0: log(f"WARNING: {trainable_vision_params:,} trainable params matched a " f"vision/projector name pattern and were forcibly frozen just now. " f"Prefixes: {sorted(matched_prefixes)}") else: log("Confirmed: zero trainable params match a vision/projector name pattern.") return result # --------------------------------------------------------------------------- # 4. Dataset loading + chat-template formatting # --------------------------------------------------------------------------- def load_and_format_dataset(dataset_path: Path, tokenizer): from datasets import load_dataset log(f"Loading dataset from {dataset_path}") ds = load_dataset("json", data_files=str(dataset_path), split="train") log(f"Loaded {len(ds)} examples") def _format(example): text = tokenizer.apply_chat_template( example["messages"], tokenize=False, add_generation_prompt=False, reasoning_strength=REASONING_STRENGTH, ) return {"text": text} ds = ds.map(_format, remove_columns=[c for c in ds.column_names if c != "messages"]) log("Formatted dataset with tokenizer.apply_chat_template " f"(reasoning_strength='{REASONING_STRENGTH}')") log(f"Sample formatted example (first 800 chars):\n{ds[0]['text'][:800]}") return ds # --------------------------------------------------------------------------- # 5. Response-only masking # --------------------------------------------------------------------------- def apply_response_masking(trainer): """Wrap trainer.train_dataset so only assistant turns contribute to the loss, using the template's real turn markers. See briefs/task-3-report.md for the standalone (tokenizer-only) verification that this masks correctly across a 2-turn and a multi-turn example before this was ever wired into a trainer.""" from unsloth_zoo.dataset_utils import train_on_responses_only log(f"Applying train_on_responses_only: instruction_part={INSTRUCTION_PART!r} " f"response_part={RESPONSE_PART!r}") trainer = train_on_responses_only( trainer, instruction_part=INSTRUCTION_PART, response_part=RESPONSE_PART, ) return trainer # --------------------------------------------------------------------------- # 7-8. Pausable training # --------------------------------------------------------------------------- def make_pause_callback(sentinel_path: Path): """Builds a TrainerCallback subclass at call time (transformers is imported lazily, matching the rest of this script's import style) that polls the PAUSE sentinel file on_step_end. CAUGHT LIVE DURING THE SMOKE LEG (see report): an earlier version of this defined `class _PauseCallback(TrainerCallback, PauseCallback)` -- a plain mixin combined via multiple inheritance -- which is a real Python MRO trap: `TrainerCallback` (listed first) defines its OWN no-op `on_step_end` stub, which method resolution order finds before the real implementation further down the MRO, silently shadowing it. The callback was then a structural no-op: `trainer.add_callback(...)` succeeded, training ran, but on_step_end never actually executed the sentinel check, so touching PAUSE mid-run had no effect at all -- confirmed live (sentinel sat untouched for 2+ full steps past creation). Fixed by inheriting from TrainerCallback directly, which is also just simpler.""" from transformers import TrainerCallback class PauseCallback(TrainerCallback): """on_step_end: if the sentinel exists, force an immediate checkpoint save, log a clear PAUSED line, stop training cleanly, and delete the sentinel itself (so the next launch doesn't immediately re-pause).""" def __init__(self, sentinel_path: Path): self.sentinel_path = Path(sentinel_path) def on_step_end(self, args, state, control, **kwargs): if self.sentinel_path.exists(): log(f"PAUSE sentinel found at {self.sentinel_path} -- pausing at " f"step {state.global_step}.") control.should_save = True control.should_training_stop = True try: self.sentinel_path.unlink() log(f"Deleted PAUSE sentinel {self.sentinel_path}") except FileNotFoundError: pass log(f"PAUSED at step {state.global_step}") return control return PauseCallback(sentinel_path) # --------------------------------------------------------------------------- # 12. Held-out generation # --------------------------------------------------------------------------- def build_holdout_prompts(holdout_dir: Path) -> list[dict]: """Builds exactly the 3 required held-out prompts (translation Sigma->KQL, explanation, authoring), all sourced from dataset/holdout/*.yml so none of them can have leaked into training. Prompt phrasing mirrors scripts/build_dataset.py's own template banks so the smoke-test prompts are representative of the training distribution.""" yml_files = sorted(holdout_dir.glob("*.yml")) if len(yml_files) < 3: raise RuntimeError(f"Expected >=3 holdout files in {holdout_dir}, found {len(yml_files)}") import yaml def _load(path): raw = path.read_text(encoding="utf-8") parsed = yaml.safe_load(raw) return raw, parsed raw0, rule0 = _load(yml_files[0]) raw1, rule1 = _load(yml_files[1]) raw2, rule2 = _load(yml_files[2]) translation_prompt = ( f"Convert this Sigma rule to Microsoft 365 Defender Advanced Hunting KQL:\n\n" f"```yaml\n{raw0}\n```" ) explanation_prompt = f"Explain this Sigma rule in plain English:\n\n```yaml\n{raw1}\n```" logsource_str = ", ".join(f"{k}={v}" for k, v in (rule2.get("logsource") or {}).items()) tags_str = ", ".join(rule2.get("tags") or []) authoring_prompt = ( f"Write a Sigma rule that detects: {rule2.get('description', rule2.get('title', ''))}\n\n" f"Logsource: {logsource_str}\nRelevant ATT&CK tags: {tags_str}" ) return [ {"task": "translation_sigma_to_kql", "source_file": yml_files[0].name, "prompt": translation_prompt}, {"task": "explanation", "source_file": yml_files[1].name, "prompt": explanation_prompt}, {"task": "authoring", "source_file": yml_files[2].name, "prompt": authoring_prompt}, ] def run_holdout_generations(model, tokenizer, holdout_dir: Path, out_path: Path): from unsloth import FastLanguageModel prompts = build_holdout_prompts(holdout_dir) FastLanguageModel.for_inference(model) results = [] for item in prompts: messages = [{"role": "user", "content": item["prompt"]}] inputs = tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, reasoning_strength=REASONING_STRENGTH, return_tensors="pt", ).to(model.device) t0 = time.time() out_ids = model.generate( input_ids=inputs, max_new_tokens=512, do_sample=True, temperature=1.0, top_p=0.95, top_k=64, ) gen_text = tokenizer.decode(out_ids[0][inputs.shape[1]:], skip_special_tokens=False) elapsed = time.time() - t0 log(f"Generated for task={item['task']} in {elapsed:.1f}s") results.append({**item, "generation": gen_text, "seconds": elapsed}) out_path.write_text(json.dumps(results, indent=2), encoding="utf-8") log(f"Wrote {len(results)} held-out generations to {out_path}") return results # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--dataset_path", type=Path, default=DATASET_PATH_DEFAULT) ap.add_argument("--output_dir", type=Path, default=OUTPUT_DIR_DEFAULT) ap.add_argument("--holdout_dir", type=Path, default=HOLDOUT_DIR_DEFAULT) ap.add_argument("--pause_sentinel", type=Path, default=PAUSE_SENTINEL_DEFAULT) ap.add_argument("--max_seq_length", type=int, default=4096) # Defaults start at OOM-ladder rung (a): 22.2GB of weights alone on one 3090 # (confirmed live) makes single-GPU seq-4096 batch-2 (the brief's literal # starting point) almost certain to OOM, so this script starts one rung down # -- batch 1 / grad-accum 16 keeps the same effective batch size (16) with far # less peak activation memory. Pass --per_device_train_batch_size 2 # --gradient_accumulation_steps 8 explicitly to try the brief's literal # starting rung anyway. ap.add_argument("--per_device_train_batch_size", type=int, default=1) ap.add_argument("--gradient_accumulation_steps", type=int, default=16) ap.add_argument("--device_map", type=str, default="sequential", help="'sequential' for single-GPU (use with CUDA_VISIBLE_DEVICES=1 -- " "GPU 0 carries ~1-2GB of desktop apps, GPU 1 is clean), " "'balanced' to shard across both GPUs (OOM fallback rung (b)).") ap.add_argument("--max_steps", type=int, default=None, help="The smoke leg actually used 20 (see briefs/task-3-report.md " "for why, not the originally-planned 100) at " "--max_seq_length 2048 (~229 steps/epoch at that config). " "Omit for a full run (num_train_epochs=1) -- NOT invoked by " "this task.") ap.add_argument("--resume", action="store_true", help="Resume from the latest checkpoint in --output_dir " "(weights + optimizer + scheduler + step).") ap.add_argument("--skip_generation", action="store_true", help="Skip the post-training held-out generation step.") ap.add_argument("--save_steps", type=int, default=200, help="Checkpoint interval in optimizer steps. The full run " "uses a tight interval (e.g. 20) so an external kill " "costs at most ~40 min at rung (c) step times.") ap.add_argument("--warmup_steps", type=int, default=None, help="Explicit override for warmup step count, used only on " "installs where SFTConfig has no warmup_ratio field (see " "the warmup_ratio/warmup_steps translation note below). " "Recommended for a full run once the real packed dataset " "size is known; the smoke leg derives it from --max_steps.") args = ap.parse_args() log(f"HF_HOME={os.environ['HF_HOME']} HF_HUB_OFFLINE={os.environ['HF_HUB_OFFLINE']}") assert_cache_hit(MODEL_NAME) model, tokenizer, loader_name = load_model_and_tokenizer( max_seq_length=args.max_seq_length, device_map=args.device_map ) log(f"Loaded via {loader_name} (raw processing object type: {type(tokenizer).__name__})") # Muse-Glimmer's FastLanguageModel.from_pretrained returns a MuseGlimmerProcessor, # not a plain tokenizer -- its __call__'s first positional arg is `images`, so any # positional tokenizer(text) call downstream (TRL's packing/collator internals # included) would misparse text as image data. Use the plain inner tokenizer for # every text-only operation from here on (see get_text_tokenizer's docstring). tokenizer = get_text_tokenizer(tokenizer) log(f"Using plain text tokenizer for all downstream ops: {type(tokenizer).__name__}") model = apply_lora_with_frozen_vision(model, loader_name) freeze_stats = verify_freeze(model) ds = load_and_format_dataset(args.dataset_path, tokenizer) from trl import SFTConfig, SFTTrainer args.output_dir.mkdir(parents=True, exist_ok=True) # transformers 5.15.0 (the version the orchestrator upgraded to, to unblock # loading -- see the report) dropped the `warmup_ratio` field from # TrainingArguments entirely (confirmed: 'warmup_ratio' not in # inspect.signature(TrainingArguments.__init__).parameters; only # `warmup_steps` remains). Unsloth's own SFTConfig shim silently drops # unknown kwargs with a warning rather than erroring, so passing # warmup_ratio=0.03 as before would train with ZERO warmup and no error -- # caught live during this run (see report). Detect and translate rather than # silently losing the brief-mandated 3% warmup. import inspect as _inspect STEPS_PER_EPOCH_ESTIMATE = 115 # from the real token-count analysis in the report warmup_ratio = 0.03 supports_warmup_ratio = "warmup_ratio" in _inspect.signature(SFTConfig.__init__).parameters warmup_kwarg = {} if supports_warmup_ratio: warmup_kwarg["warmup_ratio"] = warmup_ratio log("warmup_ratio is supported natively by the installed TRL/transformers.") elif args.warmup_steps is not None: warmup_kwarg["warmup_steps"] = args.warmup_steps log(f"Using explicit --warmup_steps={args.warmup_steps} (warmup_ratio unsupported).") else: total_steps_for_warmup = args.max_steps if args.max_steps else STEPS_PER_EPOCH_ESTIMATE warmup_steps = max(1, round(warmup_ratio * total_steps_for_warmup)) warmup_kwarg["warmup_steps"] = warmup_steps import transformers as _tf log(f"WARNING: installed TrainingArguments (transformers {_tf.__version__}) has no " f"warmup_ratio field -- translating the brief's warmup_ratio=0.03 into " f"warmup_steps={warmup_steps} (3% of {total_steps_for_warmup} total steps" + ("" if args.max_steps else " [steps/epoch ESTIMATE from token-count analysis in the report -- pass " "--warmup_steps explicitly for a precise full-run value once the real packed " "dataset size is known]") + ").") sft_config = SFTConfig( output_dir=str(args.output_dir), per_device_train_batch_size=args.per_device_train_batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, learning_rate=2e-4, num_train_epochs=1, max_steps=args.max_steps if args.max_steps else -1, optim="adamw_8bit", bf16=True, gradient_checkpointing=True, lr_scheduler_type="cosine", **warmup_kwarg, seed=42, logging_steps=5, save_steps=args.save_steps, save_total_limit=3, packing=True, max_length=args.max_seq_length, dataset_text_field="text", report_to="none", ) trainer = SFTTrainer( model=model, processing_class=tokenizer, train_dataset=ds, args=sft_config, ) trainer = apply_response_masking(trainer) trainer.add_callback(make_pause_callback(args.pause_sentinel)) import torch torch.cuda.reset_peak_memory_stats() t_train_start = time.time() trainer.train(resume_from_checkpoint=True if args.resume else False) train_wall_s = time.time() - t_train_start peak_vram = torch.cuda.max_memory_allocated() / (1024 ** 3) log(f"Training loop finished/paused. Wall clock: {train_wall_s:.1f}s. " f"Peak VRAM: {peak_vram:.2f} GiB") if not args.skip_generation: run_holdout_generations( model, tokenizer, args.holdout_dir, args.output_dir / "holdout_generations.json", ) log("Done.") if __name__ == "__main__": main()