| { |
| "datasets": [ |
| { |
| "name": "chatalpaca_predictive_state", |
| "repo_id": "BRlkl/chatalpaca-multiturn-predictive-state3", |
| "source_split": "train", |
| "format": "predictive_state_messages", |
| "validation_ratio": 0.02, |
| "split_seed": 17, |
| "min_turns": 2, |
| "max_turns": 6, |
| "max_message_chars": 6000, |
| "use_base_chat_template": true |
| } |
| ], |
| "model": { |
| "base_model_name": "google/t5gemma-l-l-prefixlm-it", |
| "initial_model_path": "BRlkl/test", |
| "dtype": "bfloat16", |
| "attn_implementation": "sdpa", |
| "magicnorm_eps": 1e-06, |
| "z_slots": 256, |
| "num_time_tokens": 0, |
| "use_explicit_time_features": false, |
| "gate_attention_heads": 4, |
| "max_observation_tokens": 1024, |
| "max_decoder_tokens": 1024, |
| "thought_loop_proposal_mode": "observation_hidden_compression", |
| "preserve_observation_encoder_manifold": true, |
| "observation_encoder_use_state_context": true, |
| "latent_attention_mask_mode": "full", |
| "initial_update_gate_bias": 20.0 |
| }, |
| "training": { |
| "seed": 17, |
| "num_workers": 2, |
| "gradient_checkpointing": true, |
| "mixed_precision": "bf16", |
| "max_grad_norm": 1.0, |
| "weight_decay": 0.01, |
| "backbone_learning_rate": 5e-06, |
| "new_module_learning_rate": 0.0001, |
| "adam_beta1": 0.9, |
| "adam_beta2": 0.95, |
| "adam_epsilon": 1e-08, |
| "fused_adamw": true, |
| "freeze_gate_head": true, |
| "log_every_steps": 1, |
| "eval_every_steps": 50, |
| "checkpoint_every_steps": 5000, |
| "eval_max_batches": 16, |
| "max_train_examples": 10, |
| "max_validation_examples": 10, |
| "freeze_decoder_first_fraction": 0.0, |
| "response_loss_weight": 0.45, |
| "probe_loss_weight": 0.35, |
| "probe_hard_loss_weight": 0.15, |
| "noop_probe_loss_weight": 0.05, |
| "probe_hard_top_fraction": 0.2, |
| "max_probe_branches_per_step": 8, |
| "max_probe_noop_ticks": 4, |
| "require_probe_steps": 1, |
| "drop_llm_not_stated_probes": true, |
| "drop_llm_positive_answer_not_in_prefix": true, |
| "drop_llm_future_answer_leaks": true, |
| "validation_behavior_enabled": true, |
| "validation_behavior_exact_only": true, |
| "validation_behavior_max_batches": 4, |
| "validation_behavior_max_probe_generations": 192, |
| "validation_probe_max_new_tokens": 96, |
| "validation_probe_extra_new_tokens": 8, |
| "validation_noop_survival_threshold_ticks": 64, |
| "validation_self_fed_behavior_enabled": true, |
| "validation_behavior_self_fed_max_batches": 2, |
| "validation_behavior_self_fed_max_probe_generations": 96, |
| "wandb_train_metric_keys": [ |
| "train/loss_total", |
| "train/loss_probe", |
| "train/probe_first_token_exact_match" |
| ], |
| "wandb_validation_metric_keys": [ |
| "loss_total", |
| "validation_goal_loss", |
| "behavior/probe_exact_match", |
| "behavior/self_fed/probe_exact_match", |
| "behavior/noop_survival_exact", |
| "behavior/secret_word_exact", |
| "behavior/entity_binding_exact", |
| "behavior/correction_latest_state_accuracy", |
| "behavior/not_stated_accuracy" |
| ], |
| "checkpoint_selection_metric": "validation_goal_loss", |
| "checkpoint_selection_mode": "min", |
| "validation_probe_miss_penalty": 2.0, |
| "validation_noop_miss_penalty": 1.0, |
| "validation_self_fed_miss_penalty": 1.0, |
| "validation_not_stated_miss_penalty": 0.5, |
| "validation_entity_miss_penalty": 0.75, |
| "validation_correction_miss_penalty": 0.75, |
| "checkpoint_min_probe_exact_match": 0.5, |
| "checkpoint_min_secret_word_exact": 0.5, |
| "checkpoint_min_entity_binding_exact": 0.5, |
| "checkpoint_min_not_stated_accuracy": 0.5, |
| "checkpoint_min_correction_latest_state_accuracy": 0.5, |
| "checkpoint_min_noop_survival_exact": 0.5, |
| "checkpoint_min_self_fed_probe_accuracy": 0.4, |
| "self_feedback_start_ratio": 0.0, |
| "self_feedback_end_ratio": 0.15, |
| "self_feedback_start_fraction": 0.65, |
| "self_feedback_end_fraction": 0.95, |
| "self_feedback_seed": 17, |
| "self_feedback_max_new_tokens": 160, |
| "self_feedback_extra_new_tokens": 8, |
| "student_observation_seed": 17, |
| "student_prev_assistant_dropout_prob": 0.0, |
| "student_prev_assistant_keep_last_chars": null |
| }, |
| "phase": { |
| "micro_batch_size": 16, |
| "eval_batch_size": 16, |
| "gradient_accumulation_steps": 1, |
| "num_train_epochs": 100, |
| "warmup_ratio": 0.03, |
| "shuffle_train": true |
| }, |
| "cache": { |
| "preprocessed_root": "cache/preprocessed_pre_sft_multiturn_predictive" |
| }, |
| "paths": { |
| "run_root": "runs_pre_sft_multiturn_predictive", |
| "export_root": "exports_multiturn_predictive" |
| }, |
| "inference": { |
| "format": "predictive_state_multiturn", |
| "use_base_chat_template": true |
| }, |
| "wandb": { |
| "enabled": true, |
| "project": "samantha-pre-sft", |
| "run_name": "t5gemma2-thoughtloop-pre-sft-predictive-state" |
| }, |
| "hub": { |
| "model_repo_id": "BRlkl/test_multiturn_predictive_state_overfit", |
| "private": false |
| } |
| } |
