avbiswas / bev-train

# Training JEV networks from Qwen
Train a network that answers choice, yes/no (`noul`), and score questions using a Qwen backbone and a small attention head. Choices share the same starting position IDs, and the attention mask prevents each choice from seeing the others. The head scores the choices without depending on their order.
Watch the video: Training JEV networks from Qwen
If you find this helpful, consider supporting on Patreon — it hosts all code, projects, slides, and write-ups from the YouTube channel.
## Start here: `arch.ipynb`
Open `arch.ipynb` for the architecture walkthrough from the livestream: tokenize the question and choices separately, build their position IDs, and construct the attention mask. The notebook explains the architecture; the Python files implement training and inference.
- `tokenization.py`: prompt formatting, choice positions, and attention masks.
- `dataset.py`: expand dataset rows into typed questions and batch them.
- `network.py`: Qwen backbone, optional LoRA, and the choice-scoring head.
- `train.py`: training, gradient accumulation, validation, and checkpoint initialization.
- `logger.py`: run metrics and LoRA/head checkpoints.
- `inference.py`: evaluation and typed answers through `answer()`.
Dataset: avbiswas/bev-decision-150K. Training loads this dataset automatically.
Original livestream: Architecture walkthrough.
## Setup
Install uv, then install the locked dependencies:
 [code] 
The first run downloads the Qwen3-0.6B weights, tokenizer, and dataset. The code selects CUDA, then Apple MPS, then CPU. Full training configs are intended for a GPU; reduce the batch size for smaller devices.
For a fresh Ubuntu GPU machine, `cloud_startup_script.sh` installs the system tools, sets up uv, and checks CUDA availability:
 [code] 
## Train
For a small overfitting sanity check, use the 10-question config:
 [code] 
This config runs 30 epochs to check whether the model can learn those questions. For full training with the first 20 Qwen layers and LoRA on the last 12 kept layers (8–19):
 [code] 
That config trains the choice head and LoRA adapters, uses 1,024 state tokens, and accumulates two microbatches of 32 questions per optimizer step. `optim.grad_accum_steps` defaults to 1. Logging, validation, and saving intervals count optimizer steps. An incomplete final accumulation group is skipped.
Other configs include `full_head_only.yaml` for a frozen backbone and `full_lora.yaml` for LoRA on the last four of 20 kept layers.
Each run writes its config, metrics, summary, and checkpoints under `runs/ - /`. Checkpoints contain the head and, when enabled, LoRA weights; the Qwen base weights are loaded separately. `best` tracks test accuracy, `latest` is overwritten at each save interval, and `final` is saved at the end.
## Start from saved weights
Pass a checkpoint directory or a run ID (which selects that run's final checkpoint):
 [code] 
Use the same backbone, head architecture, and LoRA configuration as the checkpoint.