| model: |
| n_layer: 2 |
| d_model: 3072 |
| n_ctx: 512 |
| d_head: 16 |
| d_mlp: 12288 |
| vocab_size: 4096 |
| use_rms_norm: true |
| tie_embeddings: false |
| use_positional_embeddings: false |
| use_bigram_table: false |
| use_attention_sinks: true |
| activation: gelu |
| dropout: 0 |
| use_bias: true |
| use_flash_attention: true |
| gradient_checkpointing: false |
| sparsity: |
| enable_weight_sparsity: true |
| target_l0_fraction: 0.01 |
| sparsity_anneal_start_fraction: 0.5 |
| sparsity_anneal_end_fraction: 0.75 |
| anneal_type: linear |
| min_weights_per_neuron: 4 |
| enable_activation_sparsity: true |
| activation_topk_fraction: 0.2 |
| activation_sparsity_locations: attn_q,attn_k,attn_v,attn_out,mlp_neuron,mlp_out |
| activation_topk_mode: per_token |
| mlp_neuron_pre_activation: true |
| activation_topk_annealing: linear |
| activation_topk_anneal_start_fraction: 0.01 |
| activation_topk_anneal_end_fraction: 0.33 |
| optimizer: |
| optimizer_type: adamw |
| learning_rate: 0.001 |
| beta1: 0.9 |
| beta2: 0.95 |
| weight_decay: 0.001 |
| eps: 0.1 |
| enable_grad_clip: true |
| grad_clip_rms: 1 |
| warmup_fraction: 0.01 |
| enable_lr_decay: true |
| min_lr_fraction: 0 |
| lr_decay_start_fraction: 0.8 |
| lr_decay_end_fraction: 1 |
| use_sharkfin_schedule: true |
| training: |
| dataset_name: jacobcd52/simplestories-tokenized |
| dataset_split: train |
| text_column: story |
| tokenizer_name: SimpleStories/SimpleStories-1.25M |
| total_tokens: 1000000000 |
| batch_size: 32 |
| gradient_accumulation_steps: 1 |
| mixed_precision: bf16 |
| checkpoint_dir: checkpoints/ablation_pretraining/separate_normalization |
| checkpoint_every_n_steps: 0 |
| keep_n_checkpoints: 0 |
| log_every_n_steps: 10 |
| log_gradients_every_n_steps: 100 |
| log_weights_every_n_steps: 100 |
| log_sparsity_every_n_steps: 100 |
| eval_every_n_steps: 100 |
| val_split: test |
| val_holdout_fraction: 0.01 |
| val_max_batches: 20 |
| wandb_project: nodesparse_adv_pretraining |
| wandb_run_name: separate_normalization |
| wandb_entity: null |
| use_wandb: true |
| seed: 0 |
| hf_repo: asher577/separate_normalization |
| num_workers: 2 |
| resume_from_checkpoint: null |
| use_torch_compile: false |
| torch_compile_mode: default |
| torch_compile_backend: inductor |
| ablation: |
| ci_fn_type: global_reverse_residual |
| d_resid_ci_fn: 1024 |
| reader_hidden_dims: 1024,1024 |
| use_attn_transitions: false |
| attn_n_heads: 4 |
| attn_max_len: 2048 |
| ci_targets: attn_q,attn_k,attn_v,attn_out,mlp_pre,mlp_out |
| pgd_loss_frac: 0.5 |
| stochastic_loss_frac: 0.16 |
| ci_loss_frac: 0.17 |
| unmasked_loss_frac: 0.17 |
| adv_kl_coeff: 0.5 |
| pgd_init: random |
| pgd_lr: 0.4 |
| pgd_n_steps: 1 |
| pgd_use_sign_grad: true |
| pgd_outer_lr: 0 |
| pgd_mode: batch |
| separate_grad_normalize: true |
| separate_grad_min_frac: 0.25 |
| minimality_coeff: 6.0e-08 |
| minimality_coeff_final_frac: 1 |
| minimality_coeff_annealing: true |
| minimality_coeff_anneal_start_frac: 0.25 |
| minimality_coeff_anneal_end_frac: 0.5 |
| minimality_coeff_anneal_type: exp |
| minimality_pnorm: 2 |
| minimality_eps: 1.0e-12 |
| minimality_beta: 325 |
| minimality_p_anneal_start_frac: 0.02 |
| minimality_p_anneal_final_p: 0.4 |
| minimality_p_anneal_end_frac: 0.95 |
| minimality_beta_pnorm: 2 |
| minimality_beta_p_anneal_start_frac: 0.02 |
| minimality_beta_p_anneal_final_p: 0.4 |
| minimality_beta_p_anneal_end_frac: 0.95 |
| initial_ci: 1 |
| minimality_warmup_frac: 0.02 |
| ci_thresholding: 0 |
| binarization_coeff: 0 |
| binarization_anneal_start_frac: 0.8 |
| binarization_anneal_end_frac: 1 |
| ci_learning_rate: 0.001 |
|
|
