File size: 2,450 Bytes
e95c403 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """04_backward_training.py -- demonstrate forward + backward + optimizer step.
Trains the mixer for a small number of steps on random data with a simple MSE
objective. Verifies:
1. Autograd works end-to-end (torch.compile handles the mixer's forward math)
2. Loss decreases during training (proves gradients are meaningful)
3. Both SISO and MIMO paths train correctly
Run:
python 04_backward_training.py
"""
import torch
import torch.nn.functional as F
from _loader import load_kernel
def train_one(mamba3, mode: str, n_steps: int = 10):
print(f"\n--- Training {mode.upper()} mixer for {n_steps} steps ---")
if mode == "siso":
kwargs = dict(d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1)
else:
kwargs = dict(d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4)
torch.manual_seed(0)
mixer = mamba3.NeuronMamba3Mixer(**kwargs).to("neuron")
optimizer = torch.optim.Adam(mixer.parameters(), lr=1e-3)
# Fixed target: a random tensor of the output shape
torch.manual_seed(42)
u = torch.randn(1, 64, 1024, device="neuron")
target = torch.randn(1, 64, 1024, device="neuron")
losses = []
for step in range(n_steps):
optimizer.zero_grad()
y, _ = mixer(u)
loss = F.mse_loss(y, target)
loss.backward()
optimizer.step()
try:
torch.neuron.synchronize()
except Exception:
pass
loss_val = loss.item()
losses.append(loss_val)
print(f" step {step:2d}: loss = {loss_val:.6f}")
initial, final = losses[0], losses[-1]
ratio = initial / final if final > 0 else float("inf")
decrease = initial - final
print(f"\n loss decrease: {decrease:.4f} ({initial:.4f} -> {final:.4f}, ratio {ratio:.2f}x)")
if final < initial:
print(f" [PASS] {mode} loss decreased over {n_steps} steps")
return True
else:
print(f" [FAIL] {mode} loss did not decrease")
return False
def main():
print("=" * 60)
print("Backward / Training demo: SISO + MIMO")
print("=" * 60)
mamba3 = load_kernel()
all_pass = True
for mode in ["siso", "mimo"]:
all_pass = train_one(mamba3, mode, n_steps=10) and all_pass
print("\n" + "=" * 60)
if all_pass:
print("=== TRAINING DEMO PASS ===")
else:
print("=== SOME MODES FAILED ===")
if __name__ == "__main__":
main()
|