| { |
| "network_kwargs": { |
| "class_name": "networks.modeling_llada.LLaDAModelLM.from_pretrained", |
| "pretrained_model_name_or_path": "models/LLaDA-8B-Instruct/", |
| "torch_dtype": "bfloat16", |
| "flash_attention": true |
| }, |
| "tokenizer_kwargs": { |
| "class_name": "transformers.AutoTokenizer.from_pretrained", |
| "pretrained_model_name_or_path": "models/LLaDA-8B-Instruct/" |
| }, |
| "critic_kwargs": { |
| "mlp_hidden": 512, |
| "policy_dim": 128, |
| "encoder_layers": 2, |
| "encoder_heads": 8, |
| "dropout": 0.05, |
| "dynamics_feature_dim": 8, |
| "posterior_feature_dim": 4, |
| "count_support": [ |
| 0, |
| 1, |
| 2 |
| ], |
| "count_max": 2, |
| "window_support": [ |
| 1, |
| 3, |
| 5 |
| ], |
| "initial_token_logit_count_set": 0.0, |
| "initial_count_logit_k0": 0.5, |
| "initial_count_logit_k1": 0.0, |
| "initial_count_logit_k2": -0.2, |
| "initial_window_logit_w1": 0.2, |
| "initial_window_logit_w3": 0.0, |
| "initial_window_logit_w5": -0.3 |
| }, |
| "infer_kwargs": { |
| "steps": 256, |
| "gen_length": 512, |
| "max_context_length": 0, |
| "prompt_max_tokens": 512, |
| "block_length": 32, |
| "block_steps": 16, |
| "temperature": 0.0, |
| "no_sample": true, |
| "cfg_scale": 0.0, |
| "mask_id": 126336, |
| "eos_id": 126081, |
| "unmask_policy": "confidence", |
| "remask_method": "count_set", |
| "remask_temperature": 0.4, |
| "lookback_blocks": 1, |
| "remask_min_age_current": 2, |
| "remask_max_age_lookback": 6, |
| "remask_max_updates": 1, |
| "hard_ban_remask_refill": false, |
| "ban_same_token_remask_refill": false, |
| "reforward_after_remask": true |
| }, |
| "data_loader_kwargs": [ |
| { |
| "class_name": "dataloaders.math.load_acdir_mixed_reference_dataset_and_reward", |
| "sources": [ |
| { |
| "local_path": "datasets/acdir_rl_raw", |
| "split": "train", |
| "source_dataset": "acdir_rl_raw" |
| } |
| ], |
| "batch_size": 4, |
| "num_workers": 8, |
| "seed": 313 |
| } |
| ], |
| "optimizer_kwargs": { |
| "class_name": "torch.optim.AdamW", |
| "lr": 1e-05, |
| "betas": [ |
| 0.9, |
| 0.999 |
| ], |
| "eps": 1e-08, |
| "weight_decay": 0.0 |
| }, |
| "lr_scheduler_kwargs": { |
| "class_name": "training.utils.lr_scheduler.constantlr" |
| }, |
| "acdir_kwargs": { |
| "rl_objective": "sft_gate_reject_only", |
| "count_support": [ |
| 0, |
| 1, |
| 2 |
| ], |
| "branch_sample_by_baseline": true, |
| "candidate_action_sources": [ |
| "sft_teacher" |
| ], |
| "sft_teacher_critic_ckpt": "runs/ck/critic-ckpt-002000.pt", |
| "sft_proposal_gate_only": true, |
| "freeze_critic_to_count_gate": true, |
| "sampled_guided_branches": 0, |
| "baseline_wrong_sampled_guided_branches": 0, |
| "baseline_right_sampled_guided_branches": 0, |
| "wrong_success_selection": "success_rate_first_cost_second", |
| "rescue_pair_selection": "success_rate_first_cost_second", |
| "wrong_prompt_ratio": 0.65, |
| "preference_beta": 2.0, |
| "deterministic_pair_weight": 1.0, |
| "success_loss_weight": 1.0, |
| "safety_loss_weight": 1.0, |
| "positive_min_success_rate": 0.15, |
| "right_damage_hard_negative_weight": 4.0, |
| "right_damage_margin": 0.2, |
| "top1_margin_weight": 0.0, |
| "top1_margin": 0.1, |
| "top1_margin_beta": 2.0, |
| "right_noop_loss_weight": 0.0, |
| "reject_gate_beta": 2.0, |
| "reject_gate_margin": 0.05, |
| "reject_gate_r2w_hard_weight": 10.0, |
| "reject_gate_r2w_normal_weight": 2.0, |
| "reject_gate_w2r_weight": 1.5, |
| "reject_gate_w2r_weak_weight": 0.0, |
| "reject_gate_r2r_anchor_weight": 1.0, |
| "reject_gate_w2w_anchor_weight": 0.3, |
| "sft_anchor_weight": 0.0, |
| "sft_logit_anchor_weight": 0.0, |
| "effective_change_weight": 1.0, |
| "same_token_after_deferred_weight": 0.0, |
| "skip_noop_with_cached_profile": false, |
| "profile_noop_audit_prob": 0.0, |
| "profile_cache_enabled": true, |
| "profile_cache_reuse": false, |
| "profile_replay_branch_selection": "all", |
| "profile_cache_dir": "", |
| "outer_loops": 1, |
| "profile_batches_per_loop": 48, |
| "critic_steps_per_outer_loop": 50, |
| "actor_refill_steps_per_outer_loop": 0, |
| "require_success_for_safety": false, |
| "sft_anchor_kwargs": { |
| "clean_ratio": 0.45, |
| "one_token_corrupt_ratio": 0.4, |
| "two_token_corrupt_ratio": 0.1, |
| "lambda_pl_nll": 1.0, |
| "lambda_token": 0.5, |
| "lambda_rank": 0.1, |
| "lambda_topk_margin": 0.1, |
| "lambda_clean_margin": 0.05 |
| } |
| }, |
| "func_name": "training.training_loop_acdir.training_loop", |
| "run_dir": "runs/ck/acdir/2026-06-06/81107-gpus32-batch4-bf16-ACDiR-SFTGate-RejectOnly-v1", |
| "resume_critic_ckpt": "runs/ck/acdir/2026-06-04/78664-gpus32-batch4-bf16-ACDiR-SFTGate-RejectOnly-v1/critic-ckpt-000050.pt", |
| "total_steps": 50, |
| "grad_accumulation": 1, |
| "max_grad_norm": 0.5, |
| "grad_value_clip": 10.0, |
| "seed": 313, |
| "step_per_tick": 5, |
| "snapshot_ticks": 10, |
| "state_dump_ticks": 0, |
| "precision": "bf16", |
| "allow_tf32": true, |
| "disable_gradient_checkpointing": false, |
| "wandb_enabled": true, |
| "wandb_ticks": 5, |
| "batch_size": 4, |
| "training_state_dir": "runs/ck/acdir/2026-06-06/STATES-81107" |
| } |
