Function run_compare

examples/cpu/cpu_training.py:193–236

Compare bnb AdamW vs torch AdamW on CPU to verify correctness.

(args)

def run_compare(args):
    """Compare bnb AdamW vs torch AdamW on CPU to verify correctness."""
    dtype = get_torch_dtype(args.dtype)
    print(f"=== Comparing bnb AdamW vs torch AdamW on CPU ({args.dtype}) ===\n")

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

    ds = prepare_data(tokenizer, args.dataset, args.max_length, num_samples=100)
    dataloader = torch.utils.data.DataLoader(
        ds,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
    )

    results = {}
    for label, make_opt in [
        ("bnb.AdamW", lambda m: bnb.optim.AdamW(m.parameters(), lr=args.lr)),
        ("torch.AdamW", lambda m: torch.optim.AdamW(m.parameters(), lr=args.lr)),
    ]:
        print(f"\n>> {label}")
        torch.manual_seed(42)
        model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype)
        optimizer = make_opt(model)
        history = train_loop(model, optimizer, dataloader, args.steps, args.log_interval)
        results[label] = history

    print(f"\n{'Step':>5} | {'bnb Loss':>10} | {'torch Loss':>11} | {'Diff':>10}")
    print("-" * 50)
    h_bnb = results["bnb.AdamW"]
    h_pt = results["torch.AdamW"]
    for i in range(0, min(len(h_bnb), len(h_pt)), max(1, args.log_interval)):
        s1, l1, _ = h_bnb[i]
        _, l2, _ = h_pt[i]
        print(f"{s1:5d} | {l1:10.4f} | {l2:11.4f} | {abs(l1 - l2):10.6f}")

    final_diff = abs(h_bnb[-1][1] - h_pt[-1][1])
    print(f"\nFinal loss difference: {final_diff:.6f}")
    if final_diff < 0.01:
        print("OK: bnb and torch AdamW produce nearly identical results on CPU.")
    else:
        print("NOTE: Some divergence detected (may grow over many steps).")


def run_with_trainer(args):