| { |
| "step": 2870, |
| "val_bpb": 1.474847862022957, |
| "model_config": { |
| "run": "gumbel_ablation/2rva_2gpt_learned_init", |
| "wandb_group": "gumbel_ablation", |
| "seed": 42, |
| "device_type": "cuda", |
| "config": null, |
| "depth": 4, |
| "aspect_ratio": 64, |
| "model_dim": 256, |
| "head_dim": 64, |
| "sequence_len": 1024, |
| "prefill_len": 0, |
| "vocab_size": 265, |
| "rva_blocks": [ |
| 0, |
| 1 |
| ], |
| "mamba_blocks": [], |
| "gdn_blocks": [], |
| "gla_blocks": [], |
| "moe_blocks": [], |
| "kla_blocks": [], |
| "recurrent_vocab_sizes": [ |
| 265, |
| 8192 |
| ], |
| "rva_balance_loss_weight": 0.0, |
| "rva_balance_loss_temperature": 1.0, |
| "rva_entropy_loss_weight": 0.0, |
| "rva_use_ema": true, |
| "rva_ema_decay": 0.99, |
| "rva_ema_epsilon": 1e-05, |
| "rva_use_random_reset": false, |
| "rva_reset_threshold": 1.0, |
| "rva_reset_prob": 1.0, |
| "rva_use_l2_norm": false, |
| "rva_use_fsq_ae": false, |
| "rva_fsq_levels": [ |
| 4, |
| 4, |
| 4, |
| 4, |
| 4, |
| 4, |
| 4 |
| ], |
| "rva_recon_loss_init": 1.0, |
| "rva_recon_loss_final": 0.01, |
| "rva_recon_anneal_fraction": 0.8, |
| "rva_n_embd": 0, |
| "rva_no_mlp": false, |
| "rva_unified_embed": false, |
| "rva_embed_in_muon": true, |
| "rva_learned_init_state": true, |
| "rva_pq_groups": 0, |
| "rva_pq_k_sub": 0, |
| "rva_pq_balance_loss_weight": 0.0, |
| "rva_use_gumbel": true, |
| "rva_gumbel_temp": 2.0, |
| "rva_gumbel_temp_min": 0.1, |
| "rva_gumbel_anneal_steps": 0, |
| "rva_gumbel_diversity_weight": 0.001, |
| "rva_factored": false, |
| "rva_sub_vocab_size": 0, |
| "use_skip_connection": true, |
| "clustering_init_num_samples": 0, |
| "max_active_tokens_in_batch": -1, |
| "max_active_token_fraction": 0.0, |
| "d_state": 16, |
| "mamba_params": true, |
| "kla_kernel": true, |
| "mimo_rank": 1, |
| "skip_around_kla": true, |
| "decoder_mlp": false, |
| "use_reparametrisation_trick": true, |
| "num_experts": 4, |
| "moe_top_k": 2, |
| "moe_capacity_factor": 1.0, |
| "moe_aux_loss_weight": 0.01, |
| "moe_embedding_experts": false, |
| "num_iterations": 2870, |
| "target_flops": -1.0, |
| "target_param_data_ratio": 20, |
| "data_dir": "base_data", |
| "device_batch_size": 64, |
| "total_batch_size": 524288, |
| "use_muon": true, |
| "embedding_lr": 0.3, |
| "unembedding_lr": 0.004, |
| "grad_clip": 1.0, |
| "weight_decay": 0.0, |
| "matrix_lr": 0.02, |
| "adam_beta1": 0.8, |
| "adam_beta2": 0.95, |
| "warmup_ratio": 0.05, |
| "warmdown_ratio": 0.4, |
| "final_lr_frac": 0.0, |
| "resume_from_step": -1, |
| "eval_every": 250, |
| "eval_tokens": 10485760, |
| "core_metric_every": -1, |
| "core_metric_max_per_task": 500, |
| "sample_every": 250, |
| "save_every": 1000, |
| "push_checkpoints_to_hub": true, |
| "profile_rva_codebook_usage": true, |
| "model_tag": "gumbel_ablation/2rva_2gpt_learned_init", |
| "stage": null, |
| "n_layer": 4, |
| "n_head": 4, |
| "n_kv_head": 4, |
| "n_embd": 256 |
| }, |
| "user_config": { |
| "run": "gumbel_ablation/2rva_2gpt_learned_init", |
| "wandb_group": "gumbel_ablation", |
| "seed": 42, |
| "device_type": "cuda", |
| "config": null, |
| "depth": 4, |
| "aspect_ratio": 64, |
| "model_dim": 256, |
| "head_dim": 64, |
| "sequence_len": 1024, |
| "prefill_len": 0, |
| "vocab_size": 265, |
| "rva_blocks": [ |
| 0, |
| 1 |
| ], |
| "mamba_blocks": [], |
| "gdn_blocks": [], |
| "gla_blocks": [], |
| "moe_blocks": [], |
| "kla_blocks": [], |
| "recurrent_vocab_sizes": [ |
| 265, |
| 8192 |
| ], |
| "rva_balance_loss_weight": 0.0, |
| "rva_balance_loss_temperature": 1.0, |
| "rva_entropy_loss_weight": 0.0, |
| "rva_use_ema": true, |
| "rva_ema_decay": 0.99, |
| "rva_ema_epsilon": 1e-05, |
| "rva_use_random_reset": false, |
| "rva_reset_threshold": 1.0, |
| "rva_reset_prob": 1.0, |
| "rva_use_l2_norm": false, |
| "rva_use_fsq_ae": false, |
| "rva_fsq_levels": [ |
| 4, |
| 4, |
| 4, |
| 4, |
| 4, |
| 4, |
| 4 |
| ], |
| "rva_recon_loss_init": 1.0, |
| "rva_recon_loss_final": 0.01, |
| "rva_recon_anneal_fraction": 0.8, |
| "rva_n_embd": 0, |
| "rva_no_mlp": false, |
| "rva_unified_embed": false, |
| "rva_embed_in_muon": true, |
| "rva_learned_init_state": true, |
| "rva_pq_groups": 0, |
| "rva_pq_k_sub": 0, |
| "rva_pq_balance_loss_weight": 0.0, |
| "rva_use_gumbel": true, |
| "rva_gumbel_temp": 2.0, |
| "rva_gumbel_temp_min": 0.1, |
| "rva_gumbel_anneal_steps": 0, |
| "rva_gumbel_diversity_weight": 0.001, |
| "rva_factored": false, |
| "rva_sub_vocab_size": 0, |
| "use_skip_connection": true, |
| "clustering_init_num_samples": 0, |
| "max_active_tokens_in_batch": -1, |
| "max_active_token_fraction": 0.0, |
| "d_state": 16, |
| "mamba_params": true, |
| "kla_kernel": true, |
| "mimo_rank": 1, |
| "skip_around_kla": true, |
| "decoder_mlp": false, |
| "use_reparametrisation_trick": true, |
| "num_experts": 4, |
| "moe_top_k": 2, |
| "moe_capacity_factor": 1.0, |
| "moe_aux_loss_weight": 0.01, |
| "moe_embedding_experts": false, |
| "num_iterations": 2870, |
| "target_flops": -1.0, |
| "target_param_data_ratio": 20, |
| "data_dir": "base_data", |
| "device_batch_size": 64, |
| "total_batch_size": 524288, |
| "use_muon": true, |
| "embedding_lr": 0.3, |
| "unembedding_lr": 0.004, |
| "grad_clip": 1.0, |
| "weight_decay": 0.0, |
| "matrix_lr": 0.02, |
| "adam_beta1": 0.8, |
| "adam_beta2": 0.95, |
| "warmup_ratio": 0.05, |
| "warmdown_ratio": 0.4, |
| "final_lr_frac": 0.0, |
| "resume_from_step": -1, |
| "eval_every": 250, |
| "eval_tokens": 10485760, |
| "core_metric_every": -1, |
| "core_metric_max_per_task": 500, |
| "sample_every": 250, |
| "save_every": 1000, |
| "push_checkpoints_to_hub": true, |
| "profile_rva_codebook_usage": true, |
| "model_tag": "gumbel_ablation/2rva_2gpt_learned_init", |
| "stage": null |
| }, |
| "stage": null, |
| "device_batch_size": 64, |
| "sequence_len": 1024, |
| "dataloader_state_dict": { |
| "pq_idx": 5, |
| "rg_idx": 32 |
| }, |
| "loop_state": { |
| "min_val_bpb": 1.474847862022957, |
| "smooth_train_loss": 1.0218834576418085, |
| "smooth_aux_loss": -1.963342962342586e-08, |
| "total_training_time": 907.5541379451752 |
| } |
| } |
