Skip to content

Latest commit

 

History

History
168 lines (126 loc) · 5.29 KB

File metadata and controls

168 lines (126 loc) · 5.29 KB

Fine-Tuning CLI


The voice finetune command group runs LoRA fine-tuning jobs via axolotl. Two subcommands are available: single for a single training run and sweep for a hyperparameter grid search. Both work for SFT and GRPO runs - the training mode is controlled by the config.


Single Run

voice finetune single <config> [OPTIONS]
Option Default Description
--runs-dir runs Root directory for run outputs
--gpus 0 Comma-separated CUDA device indices, e.g. 0,1,2
--dataset-revision main HuggingFace dataset branch or revision

The config file has a single top-level axolotl: key containing all training settings. The CLI injects output_dir automatically.

axolotl:
  base_model: meta-llama/Llama-3.1-8B-Instruct

  datasets:
    - path: AccelerateScience/bo-press-conference-qa
      type: chat_template
      field_messages: messages
      roles_to_train: [assistant]
      train_on_eos: last

  adapter: lora
  lora_r: 16
  lora_alpha: 16
  lora_target_modules: [q_proj, k_proj, v_proj, o_proj]

  num_epochs: 3
  learning_rate: 1.0e-4
  micro_batch_size: 4
  gradient_accumulation_steps: 4

  plugins:
    - voice.finetune.callbacks.EvalCompletionsPlugin

A full example is at configs/single/example.yaml.

When a run starts, the CLI generates a unique run name (model slug + short UUID), writes a frozen copy of the resolved config to runs/{run_name}/config.yaml, then launches axolotl. The EvalCompletionsPlugin generates completions against the validation and test splits at the end of every epoch; alignment scores are computed immediately after.

For multiple GPUs the CLI switches automatically to accelerate launch --num_processes=N.


Sweep

voice finetune sweep <config> [OPTIONS]
Option Default Description
--runs-dir runs Root directory for run outputs
--gpus 0 Comma-separated CUDA device indices
--dataset-revision main HuggingFace dataset branch or revision
--resume / --no-resume off Skip runs already recorded as completed in state.json

The sweep config has two top-level keys: sweep: defines the grid axes and axolotl: defines settings shared across every run.

sweep:
  learning_rate: [1.0e-4, 5.0e-4, 1.0e-3]
  lora_r: [4, 8, 16, 32]
  micro_batch_size: [4]
  gradient_accumulation_steps: [1, 4]
  target_layers:
    - name: attention
      modules: [q_proj, k_proj, v_proj, o_proj]
    - name: mlp
      modules: [gate_proj, up_proj, down_proj]

axolotl:
  base_model: meta-llama/Llama-3.1-8B-Instruct
  # ... shared settings ...

All five sweep: axes are required. The CLI expands them into a Cartesian product; lora_alpha is set equal to lora_r for each run and does not need to be listed. Per-run values for learning_rate, lora_r, lora_alpha, lora_target_modules, micro_batch_size, and gradient_accumulation_steps are merged on top of the shared axolotl: block before training starts.

Additional axes (such as trl.beta for GRPO sweeps) can be added freely and are passed through to the axolotl config using dot notation.

A full SFT example is at configs/sweep/example.yaml. A full GRPO example is at configs/sweep/obama/bo_qwen3_14b_grpo.yaml.

Resume

The sweep writes a state.json to {runs-dir}/state.json after each run completes or fails:

{
  "completed": ["0", "1", "3"],
  "failed": {"2": "CUDA out of memory"}
}

Passing --resume skips any run whose index already appears in completed. Failed runs are retried.


Output Layout

Each run produces a directory under {runs-dir}/{run_name}/:

runs/{run_name}/
  config.yaml            # frozen per-run axolotl config
  metadata.json          # status, hyperparams, scores
  adapter/               # axolotl output_dir; LoRA weights and checkpoints
  completions/
    epoch_1/
      validation.jsonl
      test.jsonl
    epoch_N/...
  alignment/
    epoch_1/
      validation.json    # alignment scores + confidence intervals
      test.json
    epoch_N/...

metadata.json tracks the full lifecycle of a run:

{
  "run_name": "llama-3-1-8b-instruct_a3f2b9c1",
  "model": "meta-llama/Llama-3.1-8B-Instruct",
  "dataset": "AccelerateScience/bo-press-conference-qa",
  "hyperparams": { "learning_rate": 1e-4, "lora_r": 16, "..." : "..." },
  "status": "done",
  "started_at": "2025-01-15T10:30:00Z",
  "finished_at": "2025-01-15T14:22:00Z",
  "eval_loss": { "1": 1.23, "2": 0.98, "3": 0.81 },
  "alignment_val": { "1": 0.42, "2": 0.61, "3": 0.67 },
  "alignment_test": { "1": 0.39, "2": 0.58, "3": 0.64 }
}

Weights & Biases

Add three keys anywhere inside the axolotl: block to enable W&B logging:

axolotl:
  use_wandb: true
  wandb_project: VOICE
  wandb_entity: accelerate-science

Axolotl handles all logging; no additional code is required. The run name is derived from the model slug and a unique ID and is set automatically. Omitting these keys (or setting use_wandb: false) disables W&B entirely.

The environment must be authenticated before launching a run:

wandb login
# or
export WANDB_API_KEY=<your-key>