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:

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 <think>\n (reasoning) or <think></think> (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.