| """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) |
|
|
| |
| 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() |
|
|
