File size: 4,307 Bytes
5b97bd2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
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'