Nemotron 3 5 Lightning Text2sql Lora

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.