""" Full fine-tune (not LoRA) of h2oai/h2o-danube3-500m-chat for email triage. VARIANT EXPERIMENT -- candidate replacement for cipher-nano. SmolLM2 (49K vocab, 135M/360M) shrinks to the right disk size but plateaus at weak category/importance accuracy after multiple tuning attempts (LoRA vs full-FT, epoch sweeps, data reshaping) -- a base-pretraining-quality ceiling, not a tuning problem (see DEPLOYMENT.md). Qwen2.5-0.5B has the opposite problem: strong base quality but a 151,936-token vocabulary that floors its disk size around 340-400MB regardless of quantization, so it can't shrink into nano's target range either. Danube3-500M is a plain LlamaForCausalLM with a 32,000-token vocabulary -- much smaller than Qwen/Gemma, comparable to SmolLM2 -- while coming from a more conventional larger-scale pretraining recipe (h2oai's Danube series). Worth testing whether it breaks the small-vocab-means-weak-base pattern. Same data/format/eval as the other nano candidates -- only the base model differs. Outputs: outputs/danube3-500m-full/model/ - full fine-tuned HF model Usage: python train/train_danube3_500m_full.py python train/train_danube3_500m_full.py --epochs 3 --output_dir ./my_run """ import argparse import inspect import re from pathlib import Path def parse_args(): parser = argparse.ArgumentParser(description="Full fine-tune Danube3-500M for email triage") parser.add_argument("--model_name", default="h2oai/h2o-danube3-500m-chat", help="Base HF model") parser.add_argument("--train_file", default="train.jsonl", help="Training JSONL") parser.add_argument("--val_file", default="val.jsonl", help="Validation JSONL") parser.add_argument("--output_dir", default="outputs/danube3-500m-full", help="Root output directory") parser.add_argument("--max_seq_length", type=int, default=2048) parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--lr", type=float, default=5e-5) parser.add_argument("--per_device_batch", type=int, default=2) parser.add_argument("--gradient_accumulation", type=int, default=4) parser.add_argument("--warmup_ratio", type=float, default=0.1) parser.add_argument("--seed", type=int, default=3407) parser.add_argument("--packing", action="store_true", default=False, help="Pack multiple short examples per sequence (default on)") parser.add_argument("--no-packing", dest="packing", action="store_false") return parser.parse_args() def main(args): from datasets import disable_caching, load_dataset from trl import SFTConfig, SFTTrainer from unsloth import FastLanguageModel, is_bfloat16_supported disable_caching() out_root = Path(args.output_dir) model_dir = out_root / "model" out_root.mkdir(parents=True, exist_ok=True) print(f"Loading {args.model_name} for FULL fine-tune (no LoRA, no quantization) ...") model, tokenizer = FastLanguageModel.from_pretrained( model_name=args.model_name, max_seq_length=args.max_seq_length, dtype=None, load_in_4bit=False, full_finetuning=True, ) print(f"Loading datasets: {args.train_file}, {args.val_file}") train_ds = load_dataset("json", data_files=args.train_file, split="train") val_ds = load_dataset("json", data_files=args.val_file, split="train") def format_chat(example): # Danube3-500m-chat uses its own native format -- <|prompt|>...eos # for user turns, <|answer|>...eos for assistant turns, strictly # alternating, no system role (confirmed against the tokenizer's own # chat_template and vocab: ChatML tokens aren't even present). # Fold the system prompt into the first user turn's content. msgs = example["messages"] system_content = "" if msgs and msgs[0]["role"] == "system": system_content = msgs[0]["content"] + "\n\n" msgs = msgs[1:] parts = [] first_user = True for msg in msgs: content = msg["content"] if msg["role"] == "user" and first_user: content = system_content + content first_user = False if msg["role"] == "user": parts.append(f"<|prompt|>{content.strip()}{tokenizer.eos_token}") else: parts.append(f"<|answer|>{content.strip()}{tokenizer.eos_token}") text = "".join(parts) return {"text": text} train_ds = train_ds.map(format_chat, remove_columns=train_ds.column_names) val_ds = val_ds.map(format_chat, remove_columns=val_ds.column_names) print(f"Train examples: {len(train_ds)} Validation examples: {len(val_ds)}") config_params = inspect.signature(SFTConfig).parameters training_kwargs = dict( output_dir=str(model_dir), num_train_epochs=args.epochs, per_device_train_batch_size=args.per_device_batch, per_device_eval_batch_size=args.per_device_batch, gradient_accumulation_steps=args.gradient_accumulation, learning_rate=args.lr, warmup_ratio=args.warmup_ratio, lr_scheduler_type="cosine", optim="adamw_8bit", eval_steps=100, save_strategy="steps", save_steps=100, logging_steps=10, seed=args.seed, fp16=not is_bfloat16_supported(), bf16=is_bfloat16_supported(), load_best_model_at_end=True, metric_for_best_model="eval_loss", greater_is_better=False, report_to="none", dataset_text_field="text", packing=args.packing, ) if "eval_strategy" in config_params: training_kwargs["eval_strategy"] = "steps" else: training_kwargs["evaluation_strategy"] = "steps" if "max_length" in config_params: training_kwargs["max_length"] = args.max_seq_length else: training_kwargs["max_seq_length"] = args.max_seq_length training_args = SFTConfig(**training_kwargs) trainer_kwargs = dict( model=model, train_dataset=train_ds, eval_dataset=val_ds, args=training_args, ) trainer_params = inspect.signature(SFTTrainer).parameters if "processing_class" in trainer_params: trainer_kwargs["processing_class"] = tokenizer else: trainer_kwargs["tokenizer"] = tokenizer _orig_convert_tokens_to_ids = tokenizer.convert_tokens_to_ids _sentinel_re = re.compile(r"^<([A-Z]+)_TOKEN>$") def _convert_tokens_to_ids_patched(token): match = _sentinel_re.match(token) if isinstance(token, str) else None if match: real_id = getattr(tokenizer, f"{match.group(1).lower()}_token_id", None) if real_id is not None: return real_id return _orig_convert_tokens_to_ids(token) _orig_prepare_dataset = SFTTrainer._prepare_dataset def _prepare_dataset_patched(self, dataset, processing_class, ds_args, *rest, **kw): ds_args.dataset_num_proc = None return _orig_prepare_dataset(self, dataset, processing_class, ds_args, *rest, **kw) SFTTrainer._prepare_dataset = _prepare_dataset_patched tokenizer.convert_tokens_to_ids = _convert_tokens_to_ids_patched try: trainer = SFTTrainer(**trainer_kwargs) finally: tokenizer.convert_tokens_to_ids = _orig_convert_tokens_to_ids SFTTrainer._prepare_dataset = _orig_prepare_dataset print("Starting training...") trainer.train() print(f"Saving full fine-tuned model to {model_dir}") model.save_pretrained(model_dir) tokenizer.save_pretrained(model_dir) print("Done.") if __name__ == "__main__": args = parse_args() main(args)