# FFT config for Deepseek R1 Distill Llama 3.3 70B.
#
# Requirements:
#   - Log into WandB (`wandb login`) or disable `enable_wandb`
#
# Usage:
#   oumi train -c configs/recipes/deepseek_r1/sft/distill_llama_70b/full_train.yaml
#
# See Also:
#   - Documentation: https://oumi.ai/docs/en/latest/user_guides/train/train.html
#   - Config class: oumi.core.configs.TrainingConfig
#   - Config source: https://github.com/oumi-ai/oumi/blob/main/src/oumi/core/configs/training_config.py
#   - Other training configs: configs/**/*train.yaml

model:
  model_name: "deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
  model_max_length: 2048
  torch_dtype_str: "bfloat16"
  attn_implementation: "sdpa"
  chat_template: "llama3-instruct"
  trust_remote_code: True

data:
  train:
    datasets:
      - dataset_name: "yahma/alpaca-cleaned" # 51,760 examples

training:
  trainer_type: "TRL_SFT"
  save_steps: 200
  num_train_epochs: 3
  per_device_train_batch_size: 2
  max_grad_norm: 0.0

  enable_gradient_checkpointing: True
  gradient_checkpointing_kwargs:
    use_reentrant: False
  ddp_find_unused_parameters: False
  optimizer: "adamw_torch_fused"
  learning_rate: 2.0e-05
  lr_scheduler_type: "cosine"
  warmup_ratio: 0.02

  dataloader_num_workers: "auto"
  dataloader_prefetch_factor: 16

  logging_steps: 100
  log_model_summary: False
  empty_device_cache_steps: 50
  output_dir: "output/deepseek_r1_llama70b.fft"
  include_performance_metrics: True
  enable_wandb: True

fsdp:
  enable_fsdp: True
  cpu_offload: True
  forward_prefetch: True

  sharding_strategy: "FULL_SHARD"
  state_dict_type: "SHARDED_STATE_DICT"
  auto_wrap_policy: "TRANSFORMER_BASED_WRAP"
  transformer_layer_cls: "LlamaDecoderLayer"