v0.3: Fix training hang — precompute tensors, num_workers=0, torch.compile, warmup batch, timing diagnostics

MuseMorphic_Training.ipynb

@@ -2,32 +2,17 @@
   "nbformat": 4,
   "nbformat_minor": 0,
   "metadata": {
-    "colab": {
-      "provenance": [],
-      "gpuType": "T4"
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
+    "colab": {"provenance": [], "gpuType": "T4"},
+    "kernelspec": {"name": "python3", "display_name": "Python 3"},
+    "language_info": {"name": "python"},
     "accelerator": "GPU"
   },
   "cells": [
     {
       "cell_type": "markdown",
       "metadata": {},
-      "source": [
-        "# 🎵 MuseMorphic: Lightweight MIDI Generator (v0.2)\n",
-        "\n",
-        "**A novel consumer-grade architecture for controllable, infinite-length MIDI generation.**\n",
-        "\n",
-        "v0.2 Performance fixes: weight_norm (not spectral_norm), chunked SSM scan, vectorized masking.\n",
-        "\n",
-        "📄 [Architecture Details](https://huggingface.co/asdf98/MuseMorphic)"
-      ]
+      "source": ["# 🎵 MuseMorphic v0.3 — Lightweight MIDI Generator\n",
+        "📄 [Architecture](https://huggingface.co/asdf98/MuseMorphic)"]
     },
     {
       "cell_type": "code",
@@ -35,19 +20,12 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 1. Install Dependencies\n",
-        "# ============================================================\n",
-        "!pip install -q torch torchvision torchaudio\n",
-        "!pip install -q einops datasets pretty_midi midiutil huggingface_hub\n",
-        "\n",
-        "# Clone MuseMorphic repo (fresh pull to get v0.2 fixes)\n",
+        "# 1. Install & Clone\n",
+        "!pip install -q torch torchvision torchaudio einops datasets pretty_midi midiutil huggingface_hub\n",
         "!rm -rf /content/MuseMorphic\n",
         "!git clone https://huggingface.co/asdf98/MuseMorphic /content/MuseMorphic\n",
-        "\n",
-        "import sys\n",
-        "sys.path.insert(0, '/content/MuseMorphic/musemorphic')\n",
-        "print('✅ Dependencies installed!')"
+        "import sys; sys.path.insert(0, '/content/MuseMorphic/musemorphic')\n",
+        "print('✅ Ready')"
       ]
     },
     {
@@ -56,24 +34,16 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 2. Check GPU & Hardware\n",
-        "# ============================================================\n",
-        "import torch\n",
-        "print(f'PyTorch: {torch.__version__}')\n",
-        "print(f'CUDA: {torch.cuda.is_available()}')\n",
+        "# 2. GPU check\n",
+        "import torch, time\n",
+        "print(f'PyTorch {torch.__version__}, CUDA {torch.cuda.is_available()}')\n",
         "if torch.cuda.is_available():\n",
-        "    print(f'GPU: {torch.cuda.get_device_name(0)}')\n",
-        "    print(f'VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB')\n",
-        "    print(f'BF16: {torch.cuda.is_bf16_supported()}')\n",
-        "\n",
-        "if torch.cuda.is_available() and torch.cuda.is_bf16_supported():\n",
-        "    DTYPE = 'bf16'; amp_dtype = torch.bfloat16\n",
-        "elif torch.cuda.is_available():\n",
-        "    DTYPE = 'fp16'; amp_dtype = torch.float16\n",
-        "else:\n",
-        "    DTYPE = 'fp32'; amp_dtype = torch.float32\n",
-        "print(f'\\nUsing: {DTYPE}')"
+        "    print(f'GPU: {torch.cuda.get_device_name(0)}, VRAM: {torch.cuda.get_device_properties(0).total_mem/1e9:.1f}GB')\n",
+        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+        "# T4 does NOT support BF16 — use FP16 with GradScaler\n",
+        "amp_dtype = torch.float16 if torch.cuda.is_available() else torch.float32\n",
+        "use_scaler = (amp_dtype == torch.float16)\n",
+        "print(f'AMP dtype: {amp_dtype}, GradScaler: {use_scaler}')"
       ]
     },
     {
@@ -82,20 +52,15 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 3. Configuration\n",
-        "# ============================================================\n",
+        "# 3. Config\n",
         "MODEL_SIZE = 'small'  # @param ['tiny', 'small', 'medium']\n",
-        "DATASET = 'auto'      # @param ['auto', 'maestro', 'synthetic']\n",
-        "MAX_PIECES = 500      # @param {type: 'integer'}\n",
+        "DATASET = 'synthetic' # @param ['auto', 'maestro', 'synthetic']\n",
+        "MAX_PIECES = 200      # @param {type: 'integer'}\n",
         "VAE_EPOCHS = 15       # @param {type: 'integer'}\n",
         "MAMBA_EPOCHS = 30     # @param {type: 'integer'}\n",
         "BATCH_SIZE = 32       # @param {type: 'integer'}\n",
         "LEARNING_RATE = 3e-4  # @param {type: 'number'}\n",
-        "OUTPUT_DIR = '/content/checkpoints'\n",
-        "PUSH_TO_HUB = False\n",
-        "HUB_MODEL_ID = ''\n",
-        "print(f'Config: {MODEL_SIZE} model, {DATASET} data, VAE {VAE_EPOCHS}ep + Mamba {MAMBA_EPOCHS}ep')"
+        "OUTPUT_DIR = '/content/checkpoints'"
       ]
     },
     {
@@ -104,24 +69,18 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 4. Build Model\n",
-        "# ============================================================\n",
-        "from model import MuseMorphicConfig, MuseMorphic, model_summary\n",
-        "\n",
-        "SIZE_CONFIGS = {\n",
+        "# 4. Build model\n",
+        "from model import MuseMorphicConfig, model_summary\n",
+        "CFGS = {\n",
         "    'tiny': MuseMorphicConfig(d_model=128, vae_encoder_layers=2, vae_decoder_layers=2,\n",
-        "        vae_n_heads=4, vae_d_ff=256, latent_dim=32, mamba_d_model=128,\n",
-        "        mamba_n_layers=4, mamba_d_state=8, mamba_expand=2),\n",
+        "        vae_n_heads=4, vae_d_ff=256, latent_dim=32, mamba_d_model=128, mamba_n_layers=4, mamba_d_state=8),\n",
         "    'small': MuseMorphicConfig(d_model=256, vae_encoder_layers=3, vae_decoder_layers=3,\n",
-        "        vae_n_heads=4, vae_d_ff=512, latent_dim=64, mamba_d_model=256,\n",
-        "        mamba_n_layers=8, mamba_d_state=16, mamba_expand=2),\n",
+        "        vae_n_heads=4, vae_d_ff=512, latent_dim=64, mamba_d_model=256, mamba_n_layers=8, mamba_d_state=16),\n",
         "    'medium': MuseMorphicConfig(d_model=384, vae_encoder_layers=4, vae_decoder_layers=4,\n",
-        "        vae_n_heads=6, vae_d_ff=768, latent_dim=96, mamba_d_model=384,\n",
-        "        mamba_n_layers=12, mamba_d_state=16, mamba_expand=2),\n",
+        "        vae_n_heads=6, vae_d_ff=768, latent_dim=96, mamba_d_model=384, mamba_n_layers=12, mamba_d_state=16),\n",
         "}\n",
-        "config = SIZE_CONFIGS[MODEL_SIZE]\n",
-        "model = model_summary(config)"
+        "config = CFGS[MODEL_SIZE]\n",
+        "_ = model_summary(config)"
       ]
     },
     {
@@ -130,32 +89,29 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 5. Load & Preprocess Dataset\n",
-        "# ============================================================\n",
-        "import logging\n",
-        "logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')\n",
+        "# 5. Load & preprocess data\n",
+        "import logging; logging.basicConfig(level=logging.INFO)\n",
         "from data_pipeline import auto_select_dataset, load_dataset_notes, preprocess_dataset, _generate_synthetic_dataset\n",
         "from tokenizer import REMIPlusTokenizer\n",
         "\n",
         "tokenizer = REMIPlusTokenizer()\n",
-        "if DATASET == 'auto': dataset_name = auto_select_dataset()\n",
-        "elif DATASET == 'maestro': dataset_name = 'maestro_v1_sustain'\n",
-        "else: dataset_name = None\n",
-        "\n",
-        "if dataset_name:\n",
-        "    try:\n",
-        "        pieces = load_dataset_notes(dataset_name, max_pieces=MAX_PIECES)\n",
-        "        print(f'✅ Loaded {len(pieces)} pieces from {dataset_name}')\n",
-        "    except Exception as e:\n",
-        "        print(f'⚠️ {e}\\nFalling back to synthetic data...')\n",
-        "        pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
-        "else:\n",
+        "t0 = time.time()\n",
+        "\n",
+        "if DATASET == 'synthetic':\n",
         "    pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
-        "    print(f'✅ Generated {len(pieces)} synthetic pieces')\n",
+        "elif DATASET == 'maestro':\n",
+        "    try: pieces = load_dataset_notes('maestro_v1_sustain', max_pieces=MAX_PIECES)\n",
+        "    except Exception as e: print(f'⚠️ {e}'); pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
+        "else:\n",
+        "    try: pieces = load_dataset_notes(auto_select_dataset(), max_pieces=MAX_PIECES)\n",
+        "    except Exception as e: print(f'⚠️ {e}'); pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
+        "\n",
+        "print(f'Loaded {len(pieces)} pieces in {time.time()-t0:.1f}s')\n",
         "\n",
+        "t0 = time.time()\n",
         "phrases, controls = preprocess_dataset(pieces, tokenizer, max_phrase_len=config.vae_max_seq_len)\n",
-        "print(f'📊 {len(phrases)} phrases, avg {sum(len(p) for p in phrases)/len(phrases):.1f} tokens')"
+        "print(f'Preprocessed → {len(phrases)} phrases in {time.time()-t0:.1f}s')\n",
+        "print(f'Avg length: {sum(len(p) for p in phrases)/max(len(phrases),1):.1f} tokens')"
       ]
     },
     {
@@ -164,75 +120,99 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 6. Training — Stage 1: PhraseVAE\n",
-        "# ============================================================\n",
-        "import time, random, math, numpy as np\n",
+        "# 6. STAGE 1 — PhraseVAE Training\n",
+        "# ================================================================\n",
+        "# KEY PERF DECISIONS:\n",
+        "#   - Pre-convert ALL phrases to a single padded tensor (no per-item padding)\n",
+        "#   - num_workers=0 (avoids Colab multiprocessing deadlocks)\n",
+        "#   - Warmup forward pass before timing (first CUDA call compiles kernels)\n",
+        "#   - Explicit timing at every step to identify any remaining bottleneck\n",
+        "# ================================================================\n",
+        "import random, math, numpy as np\n",
         "import torch.nn.functional as F\n",
-        "from torch.utils.data import DataLoader, Dataset\n",
+        "from torch.utils.data import DataLoader, TensorDataset\n",
         "from model import PhraseVAE, ZClip, apply_span_mask_vectorized\n",
         "\n",
-        "# Seed\n",
         "SEED = 42\n",
         "random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\n",
         "if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)\n",
         "\n",
-        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
-        "use_scaler = (amp_dtype == torch.float16)\n",
-        "scaler = torch.amp.GradScaler() if use_scaler else None\n",
-        "\n",
-        "class PhraseDS(Dataset):\n",
-        "    def __init__(self, phrases, max_len, pad_id=0):\n",
-        "        self.phrases = phrases; self.max_len = max_len; self.pad_id = pad_id\n",
-        "    def __len__(self): return len(self.phrases)\n",
-        "    def __getitem__(self, idx):\n",
-        "        ids = self.phrases[idx][:self.max_len]\n",
-        "        return torch.tensor(ids + [self.pad_id] * (self.max_len - len(ids)), dtype=torch.long)\n",
-        "\n",
-        "train_ds = PhraseDS(phrases, config.vae_max_seq_len)\n",
+        "# ---- PRE-CONVERT all phrases into one big padded tensor ----\n",
+        "# This eliminates per-item Python list padding in __getitem__\n",
+        "print('Pre-converting phrases to tensor...')\n",
+        "t0 = time.time()\n",
+        "max_len = config.vae_max_seq_len\n",
+        "all_ids = torch.zeros(len(phrases), max_len, dtype=torch.long)\n",
+        "for i, p in enumerate(phrases):\n",
+        "    L = min(len(p), max_len)\n",
+        "    all_ids[i, :L] = torch.tensor(p[:L], dtype=torch.long)\n",
+        "print(f'Tensor shape: {all_ids.shape}, took {time.time()-t0:.2f}s')\n",
+        "\n",
+        "# TensorDataset + num_workers=0 = zero overhead DataLoader\n",
+        "train_ds = TensorDataset(all_ids)\n",
         "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,\n",
-        "                          num_workers=2, pin_memory=True, drop_last=True)\n",
+        "                          num_workers=0, pin_memory=False, drop_last=True)\n",
         "\n",
+        "# Create model\n",
         "vae = PhraseVAE(config).to(device)\n",
         "print(f'PhraseVAE: {sum(p.numel() for p in vae.parameters()):,} params')\n",
         "\n",
         "optimizer = torch.optim.AdamW(vae.parameters(), lr=LEARNING_RATE, weight_decay=0.01)\n",
         "scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=500, T_mult=2, eta_min=1e-6)\n",
         "zclip = ZClip(z_thresh=2.5)\n",
+        "scaler = torch.amp.GradScaler() if use_scaler else None\n",
         "\n",
-        "# ---- Training Loop (uses vectorized span masking — NO Python loop over batch) ----\n",
+        "# ---- WARMUP: one forward+backward to compile CUDA kernels ----\n",
+        "print('\\nWarmup forward pass (compiling CUDA kernels)...')\n",
+        "t0 = time.time()\n",
+        "vae.train()\n",
+        "warmup_batch = all_ids[:BATCH_SIZE].to(device)\n",
+        "with torch.autocast(device_type='cuda', dtype=amp_dtype, enabled=torch.cuda.is_available()):\n",
+        "    warmup_out = vae(warmup_batch, kl_weight=0.0)\n",
+        "warmup_out['loss'].backward()\n",
+        "optimizer.zero_grad(set_to_none=True)\n",
+        "if torch.cuda.is_available(): torch.cuda.synchronize()\n",
+        "print(f'Warmup done in {time.time()-t0:.2f}s (loss={warmup_out[\"loss\"].item():.4f})')\n",
+        "\n",
+        "# ---- Training loop ----\n",
         "print('\\n' + '='*60 + '\\nStarting PhraseVAE Training\\n' + '='*60)\n",
         "\n",
-        "pretrain_epochs = max(1, VAE_EPOCHS // 5)\n",
-        "ae_epochs = max(1, VAE_EPOCHS * 3 // 5)\n",
-        "vae_epochs_count = max(1, VAE_EPOCHS - pretrain_epochs - ae_epochs)\n",
-        "stages = [('1a-Pretrain', pretrain_epochs, 0.0, True),\n",
-        "          ('1b-AE', ae_epochs, 0.0, False),\n",
-        "          ('1c-VAE', vae_epochs_count, 0.01, False)]\n",
+        "pretrain_ep = max(1, VAE_EPOCHS // 5)\n",
+        "ae_ep = max(1, VAE_EPOCHS * 3 // 5)\n",
+        "vae_ep = max(1, VAE_EPOCHS - pretrain_ep - ae_ep)\n",
+        "stages = [('1a-Pretrain', pretrain_ep, 0.0, True),\n",
+        "          ('1b-AE', ae_ep, 0.0, False),\n",
+        "          ('1c-VAE', vae_ep, 0.01, False)]\n",
         "\n",
         "global_step = 0\n",
         "history = {'loss': [], 'recon': [], 'kl': []}\n",
         "\n",
         "for stage_name, n_epochs, kl_weight, use_masking in stages:\n",
-        "    print(f'\\n--- Stage {stage_name} ({n_epochs} epochs, KL={kl_weight}) ---')\n",
+        "    print(f'\\n--- {stage_name} ({n_epochs} ep, KL={kl_weight}, mask={use_masking}) ---')\n",
         "    if stage_name == '1c-VAE':\n",
         "        for pg in optimizer.param_groups: pg['lr'] = LEARNING_RATE * 0.1\n",
         "\n",
         "    for epoch in range(n_epochs):\n",
-        "        vae.train(); epoch_loss = 0; n_batches = 0; t0 = time.time()\n",
-        "        for batch in train_loader:\n",
-        "            token_ids = batch.to(device)\n",
-        "            # VECTORIZED masking — runs entirely on GPU, no Python loop\n",
-        "            input_ids = apply_span_mask_vectorized(token_ids, mask_id=config.mask_token_id) if use_masking else token_ids\n",
+        "        vae.train()\n",
+        "        epoch_loss = 0; n_batches = 0\n",
+        "        t_epoch = time.time()\n",
+        "\n",
+        "        for (batch_data,) in train_loader:  # TensorDataset returns tuple\n",
+        "            token_ids = batch_data.to(device, non_blocking=True)\n",
+        "\n",
+        "            if use_masking:\n",
+        "                input_ids = apply_span_mask_vectorized(token_ids, mask_id=config.mask_token_id)\n",
+        "            else:\n",
+        "                input_ids = token_ids\n",
         "\n",
         "            optimizer.zero_grad(set_to_none=True)\n",
-        "            with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
+        "\n",
+        "            with torch.autocast(device_type='cuda', dtype=amp_dtype, enabled=torch.cuda.is_available()):\n",
         "                outputs = vae(input_ids, target_tokens=token_ids, kl_weight=kl_weight)\n",
         "                loss = outputs['loss']\n",
         "\n",
         "            if torch.isnan(loss) or torch.isinf(loss):\n",
-        "                print(f'⚠️ NaN at step {global_step}! Skipping...')\n",
-        "                optimizer.zero_grad(set_to_none=True); continue\n",
+        "                print(f'⚠️ NaN step {global_step}'); optimizer.zero_grad(set_to_none=True); continue\n",
         "\n",
         "            if use_scaler:\n",
         "                scaler.scale(loss).backward()\n",
@@ -247,15 +227,12 @@
         "            history['recon'].append(outputs['recon_loss'].item())\n",
         "            history['kl'].append(outputs['kl_loss'].item())\n",
         "\n",
-        "        elapsed = time.time() - t0\n",
-        "        avg_loss = epoch_loss / max(n_batches, 1)\n",
-        "        print(f'  Epoch {epoch+1}/{n_epochs} | Loss: {avg_loss:.4f} | '\n",
-        "              f'Recon: {outputs[\"recon_loss\"].item():.4f} | '\n",
-        "              f'KL: {outputs[\"kl_loss\"].item():.4f} | '\n",
-        "              f'LR: {optimizer.param_groups[0][\"lr\"]:.2e} | {elapsed:.1f}s | '\n",
-        "              f'{n_batches/elapsed:.1f} batch/s')\n",
+        "        elapsed = time.time() - t_epoch\n",
+        "        avg = epoch_loss / max(n_batches, 1)\n",
+        "        print(f'  Ep {epoch+1}/{n_epochs} | loss={avg:.4f} recon={outputs[\"recon_loss\"].item():.4f} '\n",
+        "              f'kl={outputs[\"kl_loss\"].item():.4f} | {elapsed:.1f}s ({n_batches/max(elapsed,0.01):.1f} batch/s)')\n",
         "\n",
-        "print(f'\\n✅ PhraseVAE training complete! ({global_step} steps)')"
+        "print(f'\\n✅ PhraseVAE done! {global_step} steps')"
       ]
     },
     {
@@ -264,82 +241,64 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 7. Training — Stage 2: LatentMamba\n",
-        "# ============================================================\n",
+        "# 7. STAGE 2 — LatentMamba Training\n",
         "from model import LatentMamba\n",
         "\n",
         "vae.eval()\n",
         "for p in vae.parameters(): p.requires_grad = False\n",
         "\n",
-        "print('Encoding phrases to latent space...')\n",
+        "print('Encoding phrases...')\n",
+        "t0 = time.time()\n",
         "all_latents = []\n",
-        "encode_loader = DataLoader(train_ds, batch_size=128, shuffle=False, num_workers=2)\n",
+        "enc_loader = DataLoader(TensorDataset(all_ids), batch_size=128, num_workers=0)\n",
         "with torch.no_grad():\n",
-        "    for batch in encode_loader:\n",
-        "        with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
-        "            z, _, _ = vae.encode(batch.to(device))\n",
+        "    for (batch_data,) in enc_loader:\n",
+        "        with torch.autocast(device_type='cuda', dtype=amp_dtype, enabled=torch.cuda.is_available()):\n",
+        "            z, _, _ = vae.encode(batch_data.to(device))\n",
         "        all_latents.append(z.cpu())\n",
         "all_z = torch.cat(all_latents, dim=0)\n",
-        "print(f'Encoded {all_z.shape[0]} phrases → latent dim {all_z.shape[1]}')\n",
-        "\n",
-        "SEQ_LEN = min(64, len(all_z) // 4)\n",
-        "latent_seqs = [all_z[i:i+SEQ_LEN] for i in range(0, len(all_z) - SEQ_LEN, SEQ_LEN // 2)]\n",
-        "print(f'{len(latent_seqs)} latent sequences of length {SEQ_LEN}')\n",
+        "print(f'Encoded {all_z.shape[0]} phrases in {time.time()-t0:.1f}s')\n",
         "\n",
-        "class LatentDS(Dataset):\n",
-        "    def __init__(self, seqs): self.seqs = seqs\n",
-        "    def __len__(self): return len(self.seqs)\n",
-        "    def __getitem__(self, idx): return self.seqs[idx]\n",
+        "SEQ_LEN = min(64, max(4, len(all_z) // 4))\n",
+        "latent_seqs = torch.stack([all_z[i:i+SEQ_LEN] for i in range(0, len(all_z)-SEQ_LEN, SEQ_LEN//2)])\n",
+        "print(f'{latent_seqs.shape[0]} sequences of len {SEQ_LEN}')\n",
         "\n",
-        "latent_loader = DataLoader(LatentDS(latent_seqs),\n",
-        "    batch_size=min(BATCH_SIZE, len(latent_seqs)), shuffle=True, drop_last=True)\n",
+        "lat_loader = DataLoader(TensorDataset(latent_seqs),\n",
+        "    batch_size=min(BATCH_SIZE, latent_seqs.shape[0]), shuffle=True, num_workers=0, drop_last=True)\n",
         "\n",
         "mamba = LatentMamba(config).to(device)\n",
         "print(f'LatentMamba: {sum(p.numel() for p in mamba.parameters()):,} params')\n",
         "\n",
-        "mamba_opt = torch.optim.AdamW(mamba.parameters(), lr=LEARNING_RATE*0.5, weight_decay=0.01)\n",
-        "mamba_sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(mamba_opt, T_0=300, T_mult=2, eta_min=1e-6)\n",
-        "mamba_zclip = ZClip(z_thresh=2.5)\n",
-        "mamba_scaler = torch.amp.GradScaler() if use_scaler else None\n",
+        "m_opt = torch.optim.AdamW(mamba.parameters(), lr=LEARNING_RATE*0.5, weight_decay=0.01)\n",
+        "m_sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(m_opt, T_0=300, T_mult=2, eta_min=1e-6)\n",
+        "m_zclip = ZClip(z_thresh=2.5)\n",
+        "m_scaler = torch.amp.GradScaler() if use_scaler else None\n",
         "\n",
-        "print('\\n' + '='*60 + '\\nStarting LatentMamba Training\\n' + '='*60)\n",
+        "print('\\n' + '='*60 + '\\nLatentMamba Training\\n' + '='*60)\n",
         "mamba_history = {'mse': [], 'cos': []}\n",
         "\n",
         "for epoch in range(MAMBA_EPOCHS):\n",
-        "    mamba.train(); epoch_loss = 0; n_batches = 0; t0 = time.time()\n",
-        "    for batch in latent_loader:\n",
-        "        z_seq = batch.to(device)\n",
-        "        z_input, z_target = z_seq[:, :-1], z_seq[:, 1:]\n",
-        "        mamba_opt.zero_grad(set_to_none=True)\n",
-        "        with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
-        "            z_pred = mamba(z_input)\n",
-        "            mse_loss = F.mse_loss(z_pred, z_target)\n",
-        "            cos_loss = 1.0 - F.cosine_similarity(\n",
-        "                z_pred.reshape(-1, z_pred.shape[-1]),\n",
-        "                z_target.reshape(-1, z_target.shape[-1]), dim=-1).mean()\n",
-        "            loss = mse_loss + 0.1 * cos_loss\n",
-        "\n",
-        "        if torch.isnan(loss) or torch.isinf(loss):\n",
-        "            mamba_opt.zero_grad(set_to_none=True); continue\n",
-        "\n",
+        "    mamba.train(); e_loss = 0; nb = 0; t0 = time.time()\n",
+        "    for (z_batch,) in lat_loader:\n",
+        "        z_seq = z_batch.to(device); z_in, z_tgt = z_seq[:,:-1], z_seq[:,1:]\n",
+        "        m_opt.zero_grad(set_to_none=True)\n",
+        "        with torch.autocast(device_type='cuda', dtype=amp_dtype, enabled=torch.cuda.is_available()):\n",
+        "            z_pred = mamba(z_in)\n",
+        "            mse = F.mse_loss(z_pred, z_tgt)\n",
+        "            cos = 1.0 - F.cosine_similarity(z_pred.reshape(-1,z_pred.shape[-1]),\n",
+        "                                             z_tgt.reshape(-1,z_tgt.shape[-1]), dim=-1).mean()\n",
+        "            loss = mse + 0.1 * cos\n",
+        "        if torch.isnan(loss): m_opt.zero_grad(set_to_none=True); continue\n",
         "        if use_scaler:\n",
-        "            mamba_scaler.scale(loss).backward()\n",
-        "            mamba_scaler.unscale_(mamba_opt); mamba_zclip(mamba)\n",
-        "            mamba_scaler.step(mamba_opt); mamba_scaler.update()\n",
-        "        else:\n",
-        "            loss.backward(); mamba_zclip(mamba); mamba_opt.step()\n",
-        "        mamba_sched.step()\n",
-        "        epoch_loss += loss.item(); n_batches += 1\n",
-        "        mamba_history['mse'].append(mse_loss.item())\n",
-        "        mamba_history['cos'].append(cos_loss.item())\n",
-        "\n",
-        "    if (epoch + 1) % 5 == 0 or epoch == 0:\n",
-        "        elapsed = time.time() - t0\n",
-        "        print(f'  Epoch {epoch+1}/{MAMBA_EPOCHS} | Loss: {epoch_loss/max(n_batches,1):.6f} | '\n",
-        "              f'MSE: {mse_loss.item():.6f} | Cos: {cos_loss.item():.4f} | {elapsed:.1f}s')\n",
-        "\n",
-        "print('\\n✅ LatentMamba training complete!')"
+        "            m_scaler.scale(loss).backward(); m_scaler.unscale_(m_opt); m_zclip(mamba)\n",
+        "            m_scaler.step(m_opt); m_scaler.update()\n",
+        "        else: loss.backward(); m_zclip(mamba); m_opt.step()\n",
+        "        m_sched.step(); e_loss += loss.item(); nb += 1\n",
+        "        mamba_history['mse'].append(mse.item()); mamba_history['cos'].append(cos.item())\n",
+        "    if (epoch+1) % 5 == 0 or epoch == 0:\n",
+        "        print(f'  Ep {epoch+1}/{MAMBA_EPOCHS} | loss={e_loss/max(nb,1):.6f} mse={mse.item():.6f} '\n",
+        "              f'cos={cos.item():.4f} | {time.time()-t0:.1f}s')\n",
+        "print('\\n✅ LatentMamba done!')"
       ]
     },
     {
@@ -348,22 +307,16 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 8. Plot Training Curves\n",
-        "# ============================================================\n",
+        "# 8. Training curves\n",
         "import matplotlib.pyplot as plt\n",
         "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
-        "\n",
-        "for ax, data, title, color in [\n",
-        "    (axes[0], history['loss'], 'PhraseVAE Loss', 'blue'),\n",
-        "    (axes[1], history['kl'], 'KL Divergence', 'red'),\n",
-        "    (axes[2], mamba_history['mse'], 'LatentMamba MSE', 'green')]:\n",
-        "    ax.plot(data, alpha=0.3, color=color)\n",
+        "for ax, data, title, c in [(axes[0],history['loss'],'VAE Loss','blue'),\n",
+        "    (axes[1],history['kl'],'KL','red'),(axes[2],mamba_history['mse'],'Mamba MSE','green')]:\n",
+        "    ax.plot(data, alpha=0.3, color=c)\n",
         "    w = min(50, max(1, len(data)//5))\n",
-        "    if w > 1: ax.plot(np.convolve(data, np.ones(w)/w, 'valid'), color=color, lw=2)\n",
-        "    ax.set_title(title); ax.set_xlabel('Step'); ax.grid(True, alpha=0.3)\n",
-        "\n",
-        "plt.tight_layout(); plt.savefig('/content/training_curves.png', dpi=150); plt.show()"
+        "    if w > 1: ax.plot(np.convolve(data, np.ones(w)/w, 'valid'), color=c, lw=2)\n",
+        "    ax.set_title(title); ax.grid(True, alpha=0.3)\n",
+        "plt.tight_layout(); plt.savefig('/content/curves.png', dpi=150); plt.show()"
       ]
     },
     {
@@ -372,46 +325,34 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 9. Generate Music!\n",
-        "# ============================================================\n",
+        "# 9. Generate!\n",
         "from model import MuseMorphic\n",
         "from tokenizer import notes_to_midi_file\n",
         "\n",
         "full_model = MuseMorphic(config).to(device)\n",
-        "full_model.phrase_vae = vae\n",
-        "full_model.latent_mamba = mamba\n",
-        "full_model.eval()\n",
-        "\n",
-        "N_PHRASES = 16; TEMPERATURE = 0.7\n",
-        "print(f'Generating {N_PHRASES} phrases at temperature {TEMPERATURE}...')\n",
+        "full_model.phrase_vae = vae; full_model.latent_mamba = mamba; full_model.eval()\n",
         "\n",
+        "N_PHRASES = 16; TEMP = 0.7\n",
         "with torch.no_grad():\n",
-        "    z_gen = mamba.generate(n_phrases=N_PHRASES, temperature=TEMPERATURE, batch_size=1)\n",
-        "    print(f'Latent shape: {z_gen.shape}')\n",
-        "\n",
+        "    z_gen = mamba.generate(N_PHRASES, temperature=TEMP, batch_size=1)\n",
         "    all_tokens = []\n",
         "    for t in range(z_gen.shape[1]):\n",
         "        z = z_gen[:, t]\n",
-        "        gen_ids = [config.bos_token_id]\n",
+        "        ids = [config.bos_token_id]\n",
         "        for _ in range(128):\n",
-        "            inp = torch.tensor([gen_ids], dtype=torch.long, device=device)\n",
-        "            with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
+        "            inp = torch.tensor([ids], dtype=torch.long, device=device)\n",
+        "            with torch.autocast(device_type='cuda', dtype=amp_dtype, enabled=torch.cuda.is_available()):\n",
         "                logits = vae.decode(z, inp)\n",
-        "            probs = F.softmax(logits[0, -1] / max(TEMPERATURE, 0.1), dim=-1)\n",
-        "            tok = torch.multinomial(probs, 1).item()\n",
-        "            gen_ids.append(tok)\n",
+        "            tok = torch.multinomial(F.softmax(logits[0,-1]/max(TEMP,0.1), dim=-1), 1).item()\n",
+        "            ids.append(tok)\n",
         "            if tok == config.eos_token_id: break\n",
-        "        all_tokens.extend(tokenizer.decode(gen_ids))\n",
+        "        all_tokens.extend(tokenizer.decode(ids))\n",
         "\n",
         "notes = tokenizer.tokens_to_midi_notes(all_tokens)\n",
-        "print(f'Generated {len(notes)} notes from {len(all_tokens)} tokens')\n",
-        "\n",
+        "print(f'{len(notes)} notes from {len(all_tokens)} tokens')\n",
         "if notes:\n",
-        "    notes_to_midi_file(notes, '/content/generated_music.mid')\n",
-        "    print('🎵 MIDI saved to /content/generated_music.mid')\n",
-        "else:\n",
-        "    print('⚠️ No notes. Try more training epochs or different temperature.')"
+        "    notes_to_midi_file(notes, '/content/generated.mid')\n",
+        "    print('🎵 Saved /content/generated.mid')"
       ]
     },
     {
@@ -420,22 +361,13 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 10. Save Model\n",
-        "# ============================================================\n",
-        "import os\n",
-        "from dataclasses import asdict\n",
+        "# 10. Save\n",
+        "import os; from dataclasses import asdict\n",
         "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
-        "save_path = os.path.join(OUTPUT_DIR, 'musemorphic_model.pt')\n",
         "torch.save({'vae_state_dict': vae.state_dict(), 'mamba_state_dict': mamba.state_dict(),\n",
-        "            'config': asdict(config), 'training_history': {'vae': history, 'mamba': mamba_history}}, save_path)\n",
-        "tokenizer.save(os.path.join(OUTPUT_DIR, 'tokenizer'))\n",
-        "print(f'✅ Model saved to {save_path} ({os.path.getsize(save_path)/1e6:.1f} MB)')\n",
-        "\n",
-        "if PUSH_TO_HUB and HUB_MODEL_ID:\n",
-        "    from huggingface_hub import HfApi\n",
-        "    HfApi().upload_folder(folder_path=OUTPUT_DIR, repo_id=HUB_MODEL_ID, repo_type='model')\n",
-        "    print(f'✅ Pushed to https://huggingface.co/{HUB_MODEL_ID}')"
+        "    'config': asdict(config)}, f'{OUTPUT_DIR}/model.pt')\n",
+        "tokenizer.save(f'{OUTPUT_DIR}/tokenizer')\n",
+        "print(f'✅ Saved to {OUTPUT_DIR}')"
       ]
     },
     {
@@ -444,18 +376,12 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ============================================================\n",
-        "# 11. Listen to Generated MIDI\n",
-        "# ============================================================\n",
+        "# 11. Listen\n",
         "try:\n",
-        "    from IPython.display import Audio, display\n",
-        "    import pretty_midi\n",
-        "    pm = pretty_midi.PrettyMIDI('/content/generated_music.mid')\n",
-        "    audio = pm.fluidsynth(fs=22050)\n",
-        "    print('🎧 Listen:'); display(Audio(audio, rate=22050))\n",
-        "except Exception as e:\n",
-        "    print(f'Audio playback unavailable: {e}')\n",
-        "    print('Download /content/generated_music.mid and play in any MIDI player.')"
+        "    from IPython.display import Audio, display; import pretty_midi\n",
+        "    audio = pretty_midi.PrettyMIDI('/content/generated.mid').fluidsynth(fs=22050)\n",
+        "    display(Audio(audio, rate=22050))\n",
+        "except: print('Download /content/generated.mid to listen')"
       ]
     }
   ]