Function run_single

examples/cpu/cpu_training.py:152–190

(args)

def run_single(args):
    dtype = get_torch_dtype(args.dtype)
    print(f"=== Training 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)
    dataloader = torch.utils.data.DataLoader(
        ds,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=collate_fn,
    )

    optimizer = create_optimizer(model, args.optimizer, args.lr)

    history = train_loop(model, optimizer, dataloader, args.steps, args.log_interval)

    loss_start = history[0][1]
    loss_end = history[-1][1]
    total_time = history[-1][2]
    print("\n--- Results ---")
    print(f"Loss: {loss_start:.4f} -> {loss_end:.4f} (delta={loss_start - loss_end:+.4f})")
    print(f"Total time: {total_time:.1f}s ({args.steps / total_time:.1f} steps/s)")
    print(f"Optimizer: bnb.optim.{args.optimizer} | Dtype: {args.dtype}")

    if loss_end >= loss_start:
        print("WARNING: Loss did not decrease! Training may not be working correctly.")
    else:
        print("OK: Loss decreased as expected.")

    return history


def run_compare(args):