Function run_with_trainer

examples/cpu/cpu_training.py:239–303

Train using HuggingFace Trainer with a bnb optimizer on CPU.

(args)

def run_with_trainer(args):
    """Train using HuggingFace Trainer with a bnb optimizer on CPU."""
    dtype = get_torch_dtype(args.dtype)
    print(f"=== Trainer mode with bnb {args.optimizer} on CPU ({args.dtype}) ===")
    print(f"Model: {args.model} | Dataset: {args.dataset}")
    print(f"Steps: {args.steps} | LR: {args.lr} | Batch: {args.batch_size} | MaxLen: {args.max_length}")
    print()

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype)

    ds = prepare_data(tokenizer, args.dataset, args.max_length)

    training_args = TrainingArguments(
        output_dir="./cpu_trainer_output",
        per_device_train_batch_size=args.batch_size,
        max_steps=args.steps,
        logging_steps=args.log_interval,
        learning_rate=args.lr,
        save_strategy="steps",
        save_steps=args.steps,
        save_total_limit=1,
        report_to="none",
        bf16=(args.dtype == "bf16"),
        use_cpu=True,
        dataloader_pin_memory=False,
    )

    optimizer = create_optimizer(model, args.optimizer, args.lr)
    scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=ds,
        data_collator=collate_fn,
        optimizers=(optimizer, scheduler),
    )

    train_result = trainer.train()
    metrics = train_result.metrics
    print("\n--- Trainer Results ---")
    print(f"Training loss: {metrics['train_loss']:.4f}")
    print(f"Training runtime: {metrics['train_runtime']:.1f}s")
    print(f"Steps/sec: {metrics['train_steps_per_second']:.1f}")
    print(f"Optimizer: bnb.optim.{args.optimizer} | Dtype: {args.dtype}")

    save_dir = "./cpu_trainer_output/final"
    print(f"\nSaving model and tokenizer to {save_dir} ...")
    trainer.save_model(save_dir)
    tokenizer.save_pretrained(save_dir)
    print("Save complete.")

    # Verify saved model can be loaded back
    print("Verifying saved model loads correctly ...")