wandb:
  entity: null
  resume: 'auto'




experiment:
    project: "rl_sdar_code" # need to be same of this file name
    function: "train"
    start_from_scratch: True
    total_step: 80
    save_every: 10
    eval_every: 10
    current_epoch: 1
    deepspeed_file: "1_node_8_gpus_deepspeed_zero3"
    num_node: 1
    node_index: 0

model:
    pretrained_model: "/abs/path/of/model" # absolute path of your model
    optimized_name: "optimized" # the output name for your optimized model, will be saved under sft_dream/ckpt
    model_base: "sdar" # set sdar for TraDo and SDAR


dataset:
    train_dataset: "PrimeIntellect" # "MATH_train""PrimeIntellect"
    optimization_data: "rl_data" # name of the rollout data output
    data_type: "code" # "math" "code" 


rollout:
    tensor_parallel_size: 1
    max_active: 256
    num_task_per_step: 64
    num_response_per_task: 16
    temperature: 1.0
    max_token: 2000
    block_size: 4
    denoising_steps_per_block: 4
    top_p: 1.0
    top_k: 0
    remasking_strategy: "low_confidence_dynamic" #"low_confidence_static""low_confidence_dynamic"
    dynamic_threshold: 0.9 # no use for "low_confidence_static"
    start_with_think: False

execute:
    num_chunk: 128 # batch size of executing codes in coding eval tasks

training:
    gradient_checkpointing_enable: True
    gradient_accumulation_steps: 32
    batch_size_lm: 1
    mixed_precision: "bf16"
    enable_tf32: True
    seed: 10086
    num_train_epochs: 1
    max_grad_norm: 1.0
    method: "TraceRL" # "random_masking" "TraceRL" "coupled"
    block_size: 4
    shrink: 1
    post_num: 0
    max_gen_length: 2000
    max_prompt_len: 784
    lower_p: 0.1
    upper_p: 0.9
    eps: 0.20
    beta: 0.01
    use_kl_estimator_k3: True


optimizer:
    name: adamw
    params: # default adamw params
        learning_rate: 1e-6
        scale_lr: False # scale learning rate by total batch size
        beta1: 0.9
        beta2: 0.999
        weight_decay: 0.0
        epsilon: 1e-8


lr_scheduler:
    scheduler: "cosine"
    params:
        learning_rate: ${optimizer.params.learning_rate}
        warmup_steps: 0
        min_lr_scale: 1.0


evaluation:
    eval_dataset: "LiveCodeBench"  # "MATH500" "LiveCodeBench"
    data_type:  "code" #"math"
    tensor_parallel_size: 1
    max_active: 256
    num_response_per_task: 3
    temperature: 1.0
    max_token: 2000
    block_size: 4
    denoising_steps_per_block: 4
    top_p: 1.0
    top_k: [0, 1]
    remasking_strategy: ["low_confidence_dynamic", "low_confidence_static"] #"low_confidence_static""low_confidence_dynamic"
    dynamic_threshold: 0.9 # no use for "low_confidence_static"
    start_with_think: False

# len(top_k) must == len(remasking_strategy)