model_kwargs:
  model_id: Llama-3.2-1B-Instruct
  sft_method: colar
  chat_template: false
  do_lora: true
  lora_config:
    r: 128
    lora_alpha: 32
  latent_cot_config:
    ce_weight: 1
    embed_modeling_weight: 1
    embed_modeling_loss: mse
    entropy_weight: 0
    pred_embed_forward_weight: 0
    max_compression_factor: 5
    pred_compressed_cot: true
    sqrt_mean: true
  latent_policy_config:
    lp_determinisitc: false
    lp_intermediate_size: 2048
  latent_generation_config:
    max_n_latent_forward: 64
    latent_temperature: 1.0
    compression_factor: 5
  answer_generation_config:
    max_new_tokens: 16
    do_sample: true
    top_p: 0.9
    temperature: 1.0
  do_rl: false
  rl_config:
    average_per_token_loss: false
    random_speed_in_group: false
    filter_dataset: false
    exp_batch_size: 8
    group_size: 8
    punish_latent_length: false
    clip_grad_norm: 1.0
    clip_eps: 0.2
    use_latent_loss: true
    use_answer_loss: true
    n_train_samples_per_epoch: 512
training_kwargs:
  optimizer:
    target: torch.optim.AdamW
    lr: 0.0001
    weight_decay: 0.01
  use_scheduler: false
  scheduler:
    target: constant_schedule_with_warmup
    warmup_steps: 1000
all_config:
  trainer:
    target: lightning.pytorch.trainer.Trainer
    devices:
    - 0
    max_steps: -1
    check_val_every_n_epoch: 5
    log_every_n_steps: 10
    num_sanity_val_steps: 2
    gradient_clip_val: 1.0
    reload_dataloaders_every_n_epochs: 0
    accumulate_grad_batches: 4
    precision: bf16-mixed
    use_distributed_sampler: true
    strategy: auto
    logger:
      target: lightning.pytorch.loggers.TensorBoardLogger
      save_dir: logs/colar
      name: qsa-coding_mix
      version: 20260802-200100_983467_coding_gsmwarm
    max_epochs: 25
  callbacks:
  - target: lightning.pytorch.callbacks.ModelCheckpoint
    save_last: true
    save_top_k: 3
    mode: max
    monitor: monitor
    auto_insert_metric_name: false
    filename: epoch{epoch}__step{step}__monitor{monitor:.3f}
    save_weights_only: true
  seed: null
  model:
    target: src.models.colar.LitCoLaR
    model_kwargs:
      model_id: Llama-3.2-1B-Instruct
      sft_method: colar
      chat_template: false
      do_lora: true
      lora_config:
        r: 128
        lora_alpha: 32
      latent_cot_config:
        ce_weight: 1
        embed_modeling_weight: 1
        embed_modeling_loss: mse
        entropy_weight: 0
        pred_embed_forward_weight: 0
        max_compression_factor: 5
        pred_compressed_cot: true
        sqrt_mean: true
      latent_policy_config:
        lp_determinisitc: false
        lp_intermediate_size: 2048
      latent_generation_config:
        max_n_latent_forward: 64
        latent_temperature: 1.0
        compression_factor: 5
      answer_generation_config:
        max_new_tokens: 16
        do_sample: true
        top_p: 0.9
        temperature: 1.0
      do_rl: false
      rl_config:
        average_per_token_loss: false
        random_speed_in_group: false
        filter_dataset: false
        exp_batch_size: 8
        group_size: 8
        punish_latent_length: false
        clip_grad_norm: 1.0
        clip_eps: 0.2
        use_latent_loss: true
        use_answer_loss: true
        n_train_samples_per_epoch: 512
    training_kwargs:
      optimizer:
        target: torch.optim.AdamW
        lr: 0.0001
        weight_decay: 0.01
      use_scheduler: false
      scheduler:
        target: constant_schedule_with_warmup
        warmup_steps: 1000
  dataloader:
    batch_size: 4
    val_batch_size: 32
    num_workers: 8
    pin_memory: true
    persistent_workers: true
  data_module:
    target: src.datasets.qsa.QSADataModule
    dataset_name: coding_mix
    tiny_dataset: false
    epoch_scaling: 1
  args:
    model: colar
    dataset: qsa
    trainer: default
    devices: '0'
    no_log: false
    log_suffix: coding_gsmwarm
    resume_ckpt_path: null
    load_ckpt_path: /content/colar_hf/logs/colar/qsa-gsm/colar-final/checkpoints/colar_best.ckpt
    workspace_path: /content/ws
    do_test: false
    test_ckpt_path: ''
    test_times: 5
    seed: 0
  unkown_args:
    dataset_name: coding_mix
    model_id: Llama-3.2-1B-Instruct
    batch_size: '4'
    accumulate_grad_batches: '4'
    max_epochs: '25'
    check_val_every_n_epoch: '5'