sft.yaml

# SFT Algorithm Configuration
sft:
  ## total number of steps to train will equal
  ## min((max_num_epochs * len(train_dataloader)), max_num_steps)
  max_num_epochs: 1
  max_num_steps: 60

  val_period: 10
  val_batches: 8
  val_global_batch_size: 32
  val_micro_batch_size: 1
  val_at_start: true
  val_at_end: false
  seed: 42

checkpointing:
  enabled: true
  checkpoint_dir: "results/sft"
  metric_name: "val:val_loss" # one of "val:" or "train:" followed by the metric name
  higher_is_better: false
  keep_top_k: 3
  save_period: 10
  checkpoint_must_save_by: null

policy:
  model_name: "meta-llama/Llama-3.2-1B"
  tokenizer:
    name: ${policy.model_name} ## specify if you'd like to use a tokenizer different from the model's default
    # chat_template can be a Jinja template string or path to a .jinja file
    chat_template: "{% for message in messages %}{%- if message['role'] == 'system'  %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user'  %}{{' Question: ' + message['content'].strip() + ' Answer:'}}{%- elif message['role'] == 'assistant'  %}{{' ' + message['content'].strip()}}{%- endif %}{% endfor %}"
    chat_template_kwargs: null # can be used to pass kwargs to the chat template, e.g., enable_thinking=true
  train_global_batch_size: 32
  train_micro_batch_size: 1
  max_total_sequence_length: 1024
  precision: "bfloat16"

  offload_optimizer_for_logprob: false

  dtensor_cfg:
    _v2: true
    enabled: true
    env_vars: {}
    cpu_offload: False
    sequence_parallel: false
    activation_checkpointing: false
    tensor_parallel_size: 1
    context_parallel_size: 1
    custom_parallel_plan: null
  
    # LoRA (Low-Rank Adaptation) Configuration
    lora_cfg:
      enabled: False  # Set to True to enable LoRA fine-tuning
      target_modules: []  # List of module names to apply LoRA (empty list with match_all_linear=true applies to all linear layers)
      exclude_modules: []  # List of module names to exclude from LoRA
      match_all_linear: true  # If True, applies LoRA to all linear layers (overrides target_modules)
      dim: 8  # LoRA rank (r): lower rank = fewer parameters but less capacity. Typical values: 4, 8, 16, 32, 64
      alpha: 32  # LoRA scaling factor: effective learning rate multiplier = alpha/dim. Typical values: 16, 32, 64
      dropout: 0.0  # Dropout probability applied to LoRA layers (0.0 = no dropout)
      dropout_position: "post"  # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA)
      lora_A_init: "xavier"  # Initialization method for LoRA A matrix: "xavier" or "uniform"
      use_triton: true  # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1

  dynamic_batching:
    enabled: false
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    sequence_length_round: 64

  sequence_packing:
    enabled: False
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    algorithm: "modified_first_fit_decreasing"
    sequence_length_round: 64

  # makes the training sequence length divisible by the tensor parallel size
  # this is useful for sequence parallel training
  make_sequence_length_divisible_by: ${policy.dtensor_cfg.tensor_parallel_size}
  max_grad_norm: 1.0

  optimizer:
    name: "torch.optim.AdamW"
    kwargs:
      lr: 5.0e-6
      weight_decay: 0.1
      betas: [0.9, 0.98]
      eps: 1e-5
      # when using Dtensor, we need to set foreach
      # and fused to False
      foreach: False
      fused: False
    
  ## ignored since enabled=false, but needed for testing purposes
  megatron_cfg:
    enabled: false
    env_vars: {}
    empty_unused_memory_level: 1
    activation_checkpointing: false
    tensor_model_parallel_size: 1
    expert_tensor_parallel_size: 1
    expert_model_parallel_size: 1
    pipeline_model_parallel_size: 1
    context_parallel_size: 1
    pipeline_dtype: ${policy.precision}
    num_layers_in_first_pipeline_stage: null
    num_layers_in_last_pipeline_stage: null
    sequence_parallel: false
    freeze_moe_router: false
    moe_router_dtype: null
    moe_router_load_balancing_type: "aux_loss"
    moe_router_bias_update_rate: 1e-3
    moe_permute_fusion: false
    #gives ~20% training perf speedup with sequence packing 
    apply_rope_fusion: True
    # gives ~25% training perf speedup with sequence packing and apply_rope_fusion
    bias_activation_fusion: True
    defer_fp32_logits: False
    moe_per_layer_logging: False
    moe_enable_deepep: false
    moe_token_dispatcher_type: "allgather"
    moe_shared_expert_overlap: false

    peft:
      enabled: false
      target_modules: []
      exclude_modules: []
      dim: 8
      alpha: 32
      dropout: 0.0
      dropout_position: "post"
      lora_A_init_method: "xavier"
      lora_B_init_method: "zero"
      a2a_experimental: false
      lora_dtype: None


    optimizer:
      optimizer: "adam" # When weight decay is set, it actually uses AdamW 
      lr: 5.0e-6
      min_lr: 4.9999e-6
      weight_decay: 0.1 # When weight decay is set, it actually uses AdamW
      bf16: true
      fp16: false
      params_dtype: "float32"

      #adam
      adam_beta1: 0.9
      adam_beta2: 0.98
      adam_eps: 1e-5

      #sgd
      sgd_momentum: 0.9

      #distributed optimizer
      use_distributed_optimizer: true
      use_precision_aware_optimizer: true

      clip_grad: ${policy.max_grad_norm}

      # optimizer cpu offload
      optimizer_cpu_offload: false
      optimizer_offload_fraction: 0.0

    scheduler:
      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
      weight_decay_incr_style: "constant"
      lr_decay_style: "constant"
      lr_decay_iters: 1000
      lr_warmup_iters: 50
      lr_warmup_init: 4.9999e-6

    distributed_data_parallel_config:
      grad_reduce_in_fp32: false
      overlap_grad_reduce: true
      overlap_param_gather: true
      data_parallel_sharding_strategy: "optim_grads_params"
      use_custom_fsdp: false

    fp8_cfg:
      enabled: false
      fp8: "e4m3"
      fp8_recipe: "blockwise"
      fp8_param: false

data:
  max_input_seq_length: ${policy.max_total_sequence_length}
  add_bos: true
  add_eos: true
  add_generation_prompt: false
  shuffle: true
  num_workers: 1

  # dataset
  train:
    dataset_name: "squad"
    split: "train"
  validation:
    dataset_name: "squad"
    split: "validation"
  # default settings for all datasets
  default:
    prompt_file: null
    system_prompt_file: null
    processor: "sft_processor"

  # You can also use multiple datasets by using a list of datasets.
  # See `examples/configs/grpo_multiple_datasets.yaml` for a full configuration example.

  # You can use custom response datasets for training and validation. For example:
  # train:
  #   # this dataset will override input_key and use the default values for other vars
  #   data_path: /path/to/local/train_dataset.jsonl
  #   input_key: question
  # validation:
  #   # this dataset will use the default values for other vars except data_path
  #   data_path: /path/to/local/val_dataset.jsonl
  # default:
  #   # will use below vars as default values if dataset doesn't specify it
  #   dataset_name: ResponseDataset
  #   input_key: input
  #   output_key: output
  #   prompt_file: null
  #   system_prompt_file: null
  #   processor: "sft_processor"
  # See https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/sft.md#datasets for more details.

  # OpenAI format specific configs
  # train_data_path: "/path/to/train.jsonl"  # Path to training data
  # val_data_path: "/path/to/val.jsonl"      # Path to validation data
  # chat_key: "messages"                     # Key for messages in the data
  # system_key: null                         # Key for system message (optional)
  # system_prompt: null                      # Default system prompt (optional)
  # tool_key: "tools"                        # Key for tools in the data
  # use_preserving_dataset: false            # If true, uses PreservingDataset to preserve heterogeneous schemas (e.g., tool calls with varying argument structures)

logger:
  log_dir: "logs"  # Base directory for all logs
  wandb_enabled: true # Make sure you do a ``wandb login [Your API key]'' before running
  tensorboard_enabled: true
  mlflow_enabled: false
  swanlab_enabled: false # Disable SwanLab logging
  monitor_gpus: true  # If true, will monitor GPU usage and log to wandb and/or tensorboard
  wandb:
    project: "sft-dev"
    name: "sft-dev-${data.train.dataset_name}"
  swanlab:
    project: "sft-dev"
    name: "sft-dev-${data.train.dataset_name}"
  tensorboard:
    log_dir: "tb_logs-sft-dev-${data.train.dataset_name}"
  mlflow:
    experiment_name: "sft-dev"
    run_name: "sft-dev-${data.train.dataset_name}"
  gpu_monitoring:
    collection_interval: 10  # How often to collect GPU usage metrics (in seconds)
    flush_interval: 10  # How often to flush GPU usage metrics to the loggers (in seconds)

cluster:
  gpus_per_node: 1
  num_nodes: 1