Kye Gomez

new examples folder

227dbb1 6 months ago

2.51 kB

import torch
from open_mythos.moda import MoDAConfig, MoDAModel


# ---------------------------------------------------------------------------
# Smoke test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    torch.manual_seed(42)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # Tiny config: 4 layers, 8 routed experts, top-2
    cfg = MoDAConfig(
        vocab_size=512,
        d_model=128,
        n_layers=4,
        n_heads_q=4,
        n_heads_kv=2,
        head_dim=32,
        max_seq_len=64,
        # MoE: 2 shared + 8 routed, activate top-2
        # (2+2)*64 = 256 ≈ equivalent to dense SwiGLU hidden~256
        n_shared_experts=2,
        n_routed_experts=8,
        n_activated_experts=2,
        expert_hidden_dim=64,
        moe_balance_alpha=0.01,
        moe_score_func="softmax",
    )

    model = MoDAModel(cfg).to(device)
    print(f"Parameters: {model.num_parameters():,}")
    print(model)

    B, T = 2, 32
    input_ids = torch.randint(0, cfg.vocab_size, (B, T), device=device)
    labels = torch.randint(0, cfg.vocab_size, (B, T), device=device)

    logits, loss = model(input_ids, labels)
    assert logits.shape == (B, T, cfg.vocab_size)
    print(f"Logits shape : {logits.shape}")
    print(f"Loss (LM + balance): {loss.item():.4f}")

    loss.backward()

    # Verify gradients
    last_writes = {
        f"blocks.{cfg.n_layers - 1}.k_write.weight",
        f"blocks.{cfg.n_layers - 1}.v_write.weight",
    }
    missing = [
        name
        for name, p in model.named_parameters()
        if p.grad is None and name not in last_writes
    ]
    if missing:
        print(f"WARNING — unexpected missing gradients: {missing}")
    else:
        print("All parameters received gradients (excluding last-block writes).")

    # Spot-check: MoE gate weights must receive gradients (through balance loss P_i)
    gate0_grad = model.blocks[0].moe.gate.weight.grad
    assert gate0_grad is not None, "blocks[0].moe.gate.weight has no gradient!"
    print(f"blocks[0].moe.gate.weight grad norm : {gate0_grad.norm().item():.6f}")

    # Spot-check: depth write projections gradient flows from layer ≥ 1 depth reads
    k0_grad = model.blocks[0].k_write.weight.grad
    assert k0_grad is not None, "blocks[0].k_write.weight has no gradient!"
    print(f"blocks[0].k_write.weight grad norm  : {k0_grad.norm().item():.6f}")

    print("Smoke test passed.")

	import torch
	from open_mythos.moda import MoDAConfig, MoDAModel


	# ---------------------------------------------------------------------------
	# Smoke test
	# ---------------------------------------------------------------------------

	if __name__ == "__main__":
	torch.manual_seed(42)
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	print(f"Device: {device}")

	# Tiny config: 4 layers, 8 routed experts, top-2
	cfg = MoDAConfig(
	vocab_size=512,
	d_model=128,
	n_layers=4,
	n_heads_q=4,
	n_heads_kv=2,
	head_dim=32,
	max_seq_len=64,
	# MoE: 2 shared + 8 routed, activate top-2
	# (2+2)*64 = 256 ≈ equivalent to dense SwiGLU hidden~256
	n_shared_experts=2,
	n_routed_experts=8,
	n_activated_experts=2,
	expert_hidden_dim=64,
	moe_balance_alpha=0.01,
	moe_score_func="softmax",
	)

	model = MoDAModel(cfg).to(device)
	print(f"Parameters: {model.num_parameters():,}")
	print(model)

	B, T = 2, 32
	input_ids = torch.randint(0, cfg.vocab_size, (B, T), device=device)
	labels = torch.randint(0, cfg.vocab_size, (B, T), device=device)

	logits, loss = model(input_ids, labels)
	assert logits.shape == (B, T, cfg.vocab_size)
	print(f"Logits shape : {logits.shape}")
	print(f"Loss (LM + balance): {loss.item():.4f}")

	loss.backward()

	# Verify gradients
	last_writes = {
	f"blocks.{cfg.n_layers - 1}.k_write.weight",
	f"blocks.{cfg.n_layers - 1}.v_write.weight",
	}
	missing = [
	name
	for name, p in model.named_parameters()
	if p.grad is None and name not in last_writes
	]
	if missing:
	print(f"WARNING — unexpected missing gradients: {missing}")
	else:
	print("All parameters received gradients (excluding last-block writes).")

	# Spot-check: MoE gate weights must receive gradients (through balance loss P_i)
	gate0_grad = model.blocks[0].moe.gate.weight.grad
	assert gate0_grad is not None, "blocks[0].moe.gate.weight has no gradient!"
	print(f"blocks[0].moe.gate.weight grad norm : {gate0_grad.norm().item():.6f}")

	# Spot-check: depth write projections gradient flows from layer ≥ 1 depth reads
	k0_grad = model.blocks[0].k_write.weight.grad
	assert k0_grad is not None, "blocks[0].k_write.weight has no gradient!"
	print(f"blocks[0].k_write.weight grad norm : {k0_grad.norm().item():.6f}")

	print("Smoke test passed.")