# Training

This guide covers training Neon VLA models — from local iteration to HuggingFace GPU Jobs.

---

## Training Architecture

Only the **action heads, fusion layer, and proprioception encoder** train. The video backbone is frozen:

```mermaid
graph TD
    subgraph "Frozen (3-7B params)"
        BB["Video Backbone<br/>Qwen2.5-Omni / Cosmos<br/>4-bit quantized"]
    end

    subgraph "Trainable (~6M params)"
        FUS["Fusion Layer"]
        PE["Proprioception Encoder"]
        AH["Action Heads<br/>Parameter Golf v2"]
        AE["Audio Encoder Projection"]
        SH["Speech Response Head"]
    end

    BB --> FUS
    PE --> FUS
    FUS --> AH
    FUS --> SH

    style BB fill:#1565c0,color:#fff
    style AH fill:#e65100,color:#fff
```

---

## Quick Start

### Local Training

```bash
python scripts/train_neon.py \
    --backbone Qwen/Qwen2.5-VL-3B-Instruct \
    --mode arms_only \
    --epochs 3 \
    --batch-size 4 \
    --dataset lerobot/xvla-agibot-world \
    --output cagataydev/neon-g1-v1
```

### HuggingFace GPU Jobs

```bash
# A100 (40GB VRAM) — recommended
hf jobs uv run \
    --flavor a100-large \
    --secrets HF_TOKEN \
    --timeout 8h \
    scripts/train_neon.py

# L4 (24GB VRAM) — 3B backbone only
hf jobs uv run \
    --flavor l4x1 \
    --secrets HF_TOKEN \
    --timeout 8h \
    scripts/train_neon.py --backbone Qwen/Qwen2.5-VL-3B-Instruct
```

---

## Configuration

### YAML Config

```yaml
# configs/g1_arms_only.yaml
model:
  backbone:
    model_id: "Qwen/Qwen2.5-VL-3B-Instruct"
    load_in_4bit: true
    freeze_backbone: true
  control_mode: "arms_only"
  include_locomotion: true
  num_action_steps: 16
  use_action_chunking: true
  use_separate_heads: true

data:
  action_chunk_size: 16
  sources:
    - name: "agibot-world"
      type: "lerobot"
      path: "lerobot/xvla-agibot-world"
      weight: 2.0

training:
  epochs: 3
  batch_size: 4
  gradient_accumulation_steps: 8
  learning_rate: 2e-4
  weight_decay: 0.01
  lr_scheduler: "cosine"
  optim: "paged_adamw_8bit"
  bf16: true
  push_to_hub: true
  hub_model_id: "cagataydev/neon-g1-v1"
```

### Python Config

```python
from neon.training.config import TrainConfig
from neon.model.neon_vla import NeonConfig
from neon.model.video_backbone import BackboneConfig
from neon.data.data_soup import DataSoupConfig, DataSourceConfig

config = TrainConfig(
    model=NeonConfig(
        backbone=BackboneConfig(
            model_id="Qwen/Qwen2.5-VL-3B-Instruct",
            load_in_4bit=True,
            freeze_backbone=True,
        ),
        control_mode="arms_only",
        num_action_steps=16,
    ),
    data=DataSoupConfig(
        sources=[
            DataSourceConfig(
                name="agibot",
                type="lerobot",
                path="lerobot/xvla-agibot-world",
                weight=2.0,
            ),
        ],
    ),
    epochs=3,
    batch_size=4,
    learning_rate=2e-4,
    push_to_hub=True,
)
```

---

## Training Loop

The `NeonTrainer` handles the full loop:

```python
from neon.training.train import NeonTrainer

trainer = NeonTrainer(config)
stats = trainer.train()

print(f"Best loss: {stats['best_loss']:.4f}")
print(f"Total steps: {stats['total_steps']}")
```

### What Happens Inside

1. **Build model** — Create NeonVLA with frozen backbone
2. **Load backbone** — Download and quantize the video model
3. **Build data loader** — Load data soup with weighted sampling
4. **Train loop** — For each batch:
   - Forward: backbone encode → fusion → action heads → predicted actions
   - Loss: MSE between predicted and target action chunks
   - Backward: gradients flow through action heads and fusion (not backbone)
   - Optimize: AdamW-8bit with gradient clipping and cosine LR schedule
5. **Checkpoint** — Save best model and upload to HuggingFace Hub

---

## Key Hyperparameters

| Parameter | Default | Notes |
|-----------|---------|-------|
| `learning_rate` | 2e-4 | Standard for action heads. Lower (1e-4) if also training LoRA on backbone |
| `batch_size` | 4 | Limited by VRAM. Each sample requires a backbone forward pass |
| `gradient_accumulation_steps` | 8 | Effective batch size = 4 × 8 = 32 |
| `num_action_steps` | 16 | Action chunk size. Match between model config and data config |
| `optim` | `paged_adamw_8bit` | Saves ~30% VRAM vs standard AdamW |
| `adam_beta1` | 0.85 | Lower than default 0.9 — faster adaptation (Parameter Golf) |
| `adam_beta2` | 0.99 | Slightly lower than default 0.999 (MicroGPT) |
| `bf16` | True | BFloat16 mixed precision. Faster and less VRAM |
| `max_grad_norm` | 0.3 | Tighter clipping for small action heads (Parameter Golf) |

---

## GPU Requirements

| Config | GPU | VRAM | Time (3 epochs, 10K samples) |
|--------|-----|------|-----|
| 3B backbone, arms only | L4 | 24 GB | ~2 hours |
| 3B backbone, whole body | RTX 4090 | 24 GB | ~3 hours |
| 7B backbone, arms only | A100 | 40 GB | ~4 hours |
| 7B backbone + LoRA | A100 80GB | 80 GB | ~8 hours |

---

## Saving and Loading

### Save

```python
model.save_pretrained("/path/to/model")
# Saves:
#   config.json          — NeonConfig
#   action_space.json    — G1ActionSpace definition
#   neon_weights.pt      — Action heads + fusion + proprio encoder weights
```

!!! note "Backbone weights are NOT saved"
    The backbone is loaded from HuggingFace at inference time. This keeps checkpoints small (~25 MB instead of 6+ GB).

### Load

```python
model = NeonVLA.from_pretrained("cagataydev/neon-g1-v1")
model.load_backbone()  # Download backbone separately
```

### Push to Hub

```python
# Automatic during training if push_to_hub=True
# Or manually:
from huggingface_hub import HfApi
api = HfApi()
api.upload_folder(
    folder_path="/path/to/model/best",
    repo_id="cagataydev/neon-g1-v1",
    repo_type="model",
)
```

---

## Lessons Learned

!!! warning "First MoE step is slow"
    If using a Mixture-of-Experts backbone, the first training step takes 10-20 minutes due to gradient checkpointing recomputing all activations. This is normal — subsequent steps are fast.

!!! warning "`flash_attention_2` not always available"
    HuggingFace GPU Jobs may not have flash-attn installed. Default to `attn_implementation="eager"` for reliability.

!!! tip "Always `push_to_hub=True`"
    HF Jobs have ephemeral storage. If the job times out, your model is lost unless it was pushed to the Hub.
