--- name: nemotron-3-5-lightning-text2sql-lora description: >- Run the Nemotron-3.5 Lightning Text2SQL LoRA fine-tuning tutorial (NeMo Megatron-Bridge) end-to-end for the user on a single node: data prep, checkpoint conversion, LoRA fine-tuning of the 30B-A3B hybrid Mamba-Transformer MoE, and merging the adapter back to a Hugging Face checkpoint. Use when the user wants to run this cookbook, fine-tune Nemotron-3.5 Lightning with LoRA, or adapt the notebook to their own machine. --- # Nemotron-3.5 Lightning Text2SQL LoRA — runbook for a coding agent This skill helps you run the cookbook in this directory (`mbridge_lora_cookbook.ipynb`) on the user's behalf. Your job is to gather a few environment details, pick a GPU configuration that fits their hardware, run the four steps in order, and confirm each one produced what it should. ## What the tutorial does Four steps, in order. **Only step 3 needs a GPU** — this is the single most useful thing to know when planning the run. 1. **Data prep** (CPU, ~2 min) — builds a BIRD Text2SQL `training.jsonl` from the no-reasoning and reasoning splits, formatted with Nemotron-3.5's chat template. 2. **Convert** (CPU, ~4 min) — imports the Hugging Face checkpoint into Megatron-Bridge format. 3. **LoRA fine-tune** (GPU) — packed-sequence LoRA training; saves an adapter. 4. **Merge & export** (CPU) — merges the adapter into the base weights, writing a standard Hugging Face checkpoint. ## Information to gather from the user Ask for these up front, in one batch: - **Path to the Nemotron-3.5 Lightning checkpoint** (already downloaded), or confirmation that you should download it and where to put it. It is ~62 GB. - **How many GPUs** they want to use, and what kind. Drives the whole training config. - **Where to write outputs** — needs ~130 GB free. - **A Hugging Face token** (`$HF_TOKEN`) so BIRD can be downloaded during data prep. Reference it by environment variable; never print it. - **The container image** to use, if it differs from the one in the notebook. ## Choosing the GPU configuration The model's 128 experts are split across GPUs with expert parallelism, so `n_devices` is the only knob that really matters — set `EP = n_devices`. Measured peak memory per GPU on 80GB H100s at `seq_length=2048`: | GPUs | Peak/GPU | One epoch | Recommendation | | --- | --- | --- | --- | | 1 | 78.8 GB | ~62 min | Works only with `REDUCE_MTP_HEADS=1`. ~0.4 GB margin — fine if that is all they have. | | 2 | 51.0 GB | ~34 min | **Default to this when available.** Stock recipe, ~28 GB margin. | | 4 | 34.8 GB | ~18 min | Good if available. | | 8 | 26.8 GB | ~8 min | Fastest. | All measured over a full epoch (189 iterations, GBS 32, `seq_length=2048`) on the complete 12,544-example dataset. Final loss lands within ~2% across all four, so choose on hardware availability and how long the user is willing to wait — not on expected quality. If the user has GPUs smaller than 80 GB, scale by the same logic: peak memory is roughly (model weights ÷ EP) + ~12 GB of overhead. Spare memory is best spent raising `seq_length`, which increases how much of the dataset survives the length filter — not just headroom. ## How to run it 1. Launch the container with the notebook directory and the checkpoint path mounted, and `$HF_TOKEN` exposed. Use the `docker run` invocation in the notebook's first cell as the template. 2. Fill in the notebook's **Configuration** cell (paths, `n_devices`) — it is the only cell that should need editing. 3. Run the steps in order. After each, run its sanity-check cell before moving on. 4. Training is the long step. Run it in the background and poll; do not hold a blocking session open, and do not stream the full log. You can also run the steps directly rather than through the notebook — each is a plain script driven by environment variables (`MODEL_ID`, `MAX_SEQ_LEN`, `DATAPREP_OUTPUT_DIR` for data prep; `HF_MODEL`, `MEGATRON_MODEL_PATH` for convert; and `N_DEVICES`, `EP`, `DATASET_DIR`, `TRAINING_OUTPUT_DIR`, `EXPERIMENT_NAME` for training). ## Verifying success per step - **Data prep:** `$DATAPREP_OUTPUT_DIR/training.jsonl` exists with ~12,500 rows at `seq_length=2048`. Spot-check one record: `input` should end with `\n` (reasoning) or `` (non-reasoning), and `output` should continue directly from there. - **Convert:** `$MEGATRON_MODEL_PATH/latest_checkpointed_iteration.txt` plus an `iter_*` directory exist (~62 GB). - **Train:** an `iter_*` adapter checkpoint under `$TRAINING_OUTPUT_DIR/$EXPERIMENT_NAME`, and the log shows `lm loss` trending down. - **Merge:** the output directory contains `model-*.safetensors` shards, `config.json`, and the tokenizer files, and the log ends with `Success: All tensors from the original checkpoint were written.` Report per-step status, wall-clock time, and the final training loss. ## Things already handled — do not change them - **The recipe supplies everything model-specific.** `train.py` calls the shipped PEFT recipe and overrides only local paths, parallelism, dataset, and schedule. Don't hand-write LoRA target modules — the recipe's already cover the Mamba projections, attention, and both routed and shared experts. - **The MoE dispatcher is set to `alltoall`** rather than the recipe's default `flex`/DeepEP, for portability. Only change this if DeepEP is known good on the user's system. - **Synchronous checkpoint saving** (`async_save=False`) is deliberate. - **Steps are idempotent:** data prep skips if `training.jsonl` exists; convert skips if the checkpoint exists. ## Expected friction (so you don't misread it) - **The first training iteration takes 1–2 minutes** with no output while CUDA graphs are captured and the MoE warms up. Subsequent iterations are seconds. Do not cancel the job. - **Startup log noise is not failure.** `Failed to import Triton kernels`, `MimoModelConfig is experimental`, `Unable to import torchao`, and `torch_dtype is deprecated` all appear on healthy runs. Judge by the sanity checks. - **Do not enable `RECOMPUTE_ACTIVATIONS`.** It lowers memory but fails at iteration 2 with an assertion in Megatron's gradient buffer. If the user is out of memory, add a GPU or lower `seq_length` instead. - **`REDUCE_MTP_HEADS` reduces to one head, it cannot disable MTP.** The hybrid model asserts `mtp_num_layers > 0`. - **If you point the recipe at local data**, you must also clear its Hugging Face dataset fields — a dataset config accepts exactly one source. `train.py` already does this; preserve it if you refactor. - **Serve with vLLM, not Transformers.** vLLM supports this architecture natively and works. `transformers.generate()` currently fails inside the model's bundled remote code — on the *base* checkpoint too, so don't diagnose it as a fine-tuning problem. - **Any vLLM script needs an `if __name__ == "__main__":` guard.** vLLM spawns workers; without it the failure surfaces as `Engine core initialization failed` wrapping a `multiprocessing` bootstrap error that never mentions vLLM. ## What success looks like at the end Serving the merged checkpoint and prompting it the way data prep formatted training examples should yield bare SQL, e.g. `SELECT T2.dept_name FROM employees AS T1 INNER JOIN ...`. The base model instead answers conversationally with fenced SQL and a prose explanation. If the fine-tuned model still explains itself, something upstream went wrong — suspect the chat-template format first.