Skip to content

posttraining.md

Latest commit

 

History

History
217 lines (152 loc) · 8.99 KB

posttraining.md

File metadata and controls

217 lines (152 loc) · 8.99 KB
Raw

Post-training workflows

Post-training (supervised fine-tuning) adapts a pre-trained foundation model to new tasks or domains. In Primus, post-training runs through the Megatron Bridge backend using the train posttrain subcommand. Example YAML configurations live under examples/megatron_bridge/configs/ in the Primus repository.

For YAML field details, see Megatron Bridge parameters. For related tooling, see Micro-benchmarking suite, Preflight diagnostics, Memory and performance projection.

Overview: SFT vs LoRA

The following table shows how SFT (Supervised Fine-Tuning) and LoRA (Low-Rank Adaptation) differ in different aspects.

Aspect SFT (full fine-tuning) LoRA (parameter-efficient)
PEFT setting peft: "none" peft: lora
What is trained All model parameters Low-rank adapters only
Memory Higher Lower
Throughput Typically slower per step Often faster iteration
Learning rate Lower: roughly 5e-6 to 1e-5 Higher: roughly 1e-4 to 5e-4
Typical use Maximum adaptation when memory allows Limited GPU memory, many task-specific adapters, rapid experimentation

Quick start commands

General form:

./primus-cli <mode> -- train posttrain --config <path-to-yaml>

From a clone of the Primus repository, the same entrypoint is often invoked as ./runner/primus-cli.

Prerequisites: AMD ROCm (recommended ≥ 7.0) and Docker with ROCm support (optional but typical) installed on systems with AMD Instinct™ GPUs (for example MI300X, MI355X). Run this for a quick check on these prerequisites: rocm-smi && docker --version. See Installation and setup for the full prerequisites and container setup.

Direct mode (bare metal or inside a Docker container)

# SFT — example: Qwen3 32B on MI355X
./runner/primus-cli direct -- train posttrain \
  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml

# LoRA — same model family
./runner/primus-cli direct -- train posttrain \
  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml

Container mode

./runner/primus-cli container --image rocm/primus:v26.8 -- \
  train posttrain \
  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml

Configuration reference

These keys are commonly set under modules.post_trainer.overrides in your configuration YAML (see the examples in the repository under examples/megatron_bridge/configs/).

Area Parameters Notes
Method peft "none" for SFT; lora for LoRA.
Learning rate finetune_lr, min_lr, lr_warmup_iters, lr_decay_iters LoRA usually needs a higher finetune_lr than SFT.
Precision precision_config Typical: bf16_mixed. Alternatives include fp16_mixed and fp32 depending on backend support.
Parallelism tensor_model_parallel_size, pipeline_model_parallel_size, context_parallel_size, sequence_parallel Increase TP/PP when the model does not fit on fewer GPUs.
Recompute (memory) recompute_granularity, recompute_method, recompute_num_layers Use to trade compute for activation memory (for example recompute_granularity: full with uniform recompute).
Batching / length train_iters, global_batch_size, micro_batch_size, seq_length micro_batch_size is per-GPU; tune with sequence length and memory.

Snippet of an SFT configuration (illustrative only):

modules:
  post_trainer:
    framework: megatron_bridge
    config: sft_trainer.yaml
    model: qwen3_32b.yaml
    overrides:
      peft: "none"
      finetune_lr: 5.0e-6
      precision_config: bf16_mixed
      tensor_model_parallel_size: 1
      global_batch_size: 8
      micro_batch_size: 1
      seq_length: 8192

Snippet of a LoRA configuration (illustrative only):

modules:
  post_trainer:
    framework: megatron_bridge
    config: sft_trainer.yaml
    model: qwen3_32b.yaml
    overrides:
      peft: lora
      finetune_lr: 1.0e-4
      precision_config: bf16_mixed
      recompute_granularity: full
      recompute_method: uniform
      recompute_num_layers: 1

The reference configurations below are organized by GPU architecture. MI325X uses the same configurations as MI300X (both are gfx942), and MI350X uses the same configurations as MI355X (both are gfx950).

MI300X configurations

Paths are relative to examples/megatron_bridge/configs/ in the Primus repository.

Model Method Config path TP GBS MBS Seq len
Qwen3 32B SFT MI300X/qwen3_32b_sft_posttrain.yaml 2 8 2 8192
Qwen3 32B LoRA MI300X/qwen3_32b_lora_posttrain.yaml 1 32 2 8192

Legend: TP = tensor parallel size; GBS = global batch size; MBS = micro batch size per GPU; Seq len = seq_length.

Sample command for running the post-training:

./runner/primus-cli direct -- train posttrain \
  --config ./examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml

MI355X configurations

Paths are relative to examples/megatron_bridge/configs/ in the Primus repository.

Model Method Config path TP GBS MBS Seq len
Qwen3 32B SFT MI355X/qwen3_32b_sft_posttrain.yaml 1 8 1 8192
Qwen3 32B LoRA MI355X/qwen3_32b_lora_posttrain.yaml 1 32 4 8192

Sample command for running the post-training:

./runner/primus-cli direct -- train posttrain \
  --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml

Best practices

Use SFT or LoRA?

  • SFT is preferred when you need the strongest possible task fit, have enough GPU memory, and can afford longer runs.
  • LoRA is preferred when memory is tight, you want fast iteration, or you plan to maintain multiple adapters for different tasks.

Learning rates

  • SFT: start in the 5e-6–1e-5 range; adjust with validation loss.
  • LoRA: often 1e-4–5e-4; still use warmup (lr_warmup_iters) for stability.

Batch sizes

  • SFT: starting with global_batch_size: 8 is a reasonable default for development; scale up when stable (for example to 64, 128, or higher) if memory and throughput allow.
  • LoRA: larger global batches are often feasible (for example 32 in the reference configs); align micro_batch_size with sequence length and available HBM.
  • Very long sequences (for example 8192) may require smaller micro-batches or more parallelism.

Parallelism

  • SFT: large models may need higher tensor_model_parallel_size (for example TP 8 for very large models). The bundled 32B examples use TP 2 on MI300X and TP 1 on MI355X for SFT.
  • LoRA: adapters reduce memory pressure; lower TP is often sufficient for a given model size.

Troubleshooting

Out of memory (OOM)

SFT

  1. Increase tensor_model_parallel_size (and/or pipeline parallelism for very large models).
  2. Reduce micro_batch_size or seq_length.
  3. Enable activation recomputation (recompute_granularity, recompute_method, recompute_num_layers).

LoRA

  1. Confirm peft: lora is set.
  2. Reduce micro_batch_size if OOM persists.
  3. Apply the same recompute settings as for SFT.

Training instability (loss spikes, NaNs)

  1. Decrease finetune_lr.
  2. Increase lr_warmup_iters.
  3. Keep mixed precision stable (precision_config: bf16_mixed where supported).
  4. Monitor gradients and clipping settings if exposed by your trainer config.

Slow training

  1. Increase effective batch size where memory allows (global_batch_size / micro_batch_size tuning).
  2. Revisit TP/PP/CP for your cluster topology.
  3. Run benchmarks or preflight to isolate network or GPU issues.

Configuration errors

  1. Verify YAML paths and indentation.
  2. Set PRIMUS_WORKSPACE and other environment variables expected by your team’s templates.
  3. Confirm checkpoint and data paths by reviewing the experiment YAML and any presets it references. The current core training runtime parses --export_config, but resolved-config export is not implemented on the default PrimusRuntime path.

Related documentation