Create config.yaml

config.yaml

@@ -0,0 +1,79 @@
+output_dir: /home/jupyter/datasphere/alignment/output  
+
+# Model Arguments
+model:
+  _component_: torchtune.models.gemma2.lora_gemma2_9b
+  lora_attn_modules: ['q_proj', 'v_proj', 'output_proj']
+  apply_lora_to_mlp: True
+  lora_rank: 8  # higher increases accuracy and memory
+  lora_alpha: 16  # usually alpha=2*rank
+  lora_dropout: 0.0
+
+# Tokenizer
+tokenizer:
+  _component_: torchtune.models.gemma.gemma_tokenizer
+  path: /home/jupyter/datasphere/alignment/epoch2/tokenizer.model
+  max_seq_len: 1024 # higher increases memory
+
+checkpointer:
+  _component_: torchtune.training.FullModelHFCheckpointer
+  checkpoint_dir: /home/jupyter/datasphere/alignment/epoch2/
+  checkpoint_files: [
+    ft-model-00001-of-00004.safetensors,
+    ft-model-00002-of-00004.safetensors,
+    ft-model-00003-of-00004.safetensors,
+    ft-model-00004-of-00004.safetensors
+  ]
+  recipe_checkpoint: null
+  output_dir: ${output_dir}
+  model_type: GEMMA2
+resume_from_checkpoint: False
+save_adapter_weights_only: False
+
+# Dataset and Sampler
+dataset:
+  _component_: torchtune.datasets.preference_dataset
+  source: json
+  data_files: /home/jupyter/datasphere/alignment/data/convo_dpo.json
+  column_map:
+    chosen: chosen_conversations
+    rejected: rejected_conversations
+  train_on_input: False
+  split: train
+seed: null
+shuffle: True
+batch_size: 4
+
+# Optimizer and Scheduler
+optimizer:
+  _component_: torch.optim.AdamW
+  fused: True
+  weight_decay: 0.05
+  lr: 5e-4
+lr_scheduler:
+  _component_: torchtune.training.lr_schedulers.get_cosine_schedule_with_warmup
+  num_warmup_steps: 100
+
+loss:
+  _component_: torchtune.rlhf.loss.SimPOLoss
+
+# Training
+epochs: 3
+max_steps_per_epoch: 1000
+gradient_accumulation_steps: 8  # Use to increase effective batch size
+compile: False  # torch.compile the model + loss, True increases speed + decreases memory
+
+# Logging
+metric_logger:
+  _component_: torchtune.training.metric_logging.DiskLogger
+  log_dir: ${output_dir}/logs
+log_every_n_steps: 1
+log_peak_memory_stats: True
+
+# Environment
+device: cuda
+dtype: bf16
+
+# Memory management
+enable_activation_checkpointing: True  # True reduces memory
+enable_activation_offloading: False  # True reduces memory