{
  "step": 19000,
  "config": {
    "backend": "tpu",
    "data": {
      "train_split": "/mnt/data/splits/train.jsonl",
      "val_split": "/mnt/data/splits/val.jsonl",
      "encoded_dir": "/mnt/data/encoded",
      "max_frames": 300,
      "bucket_frames": null,
      "audio_frame_rate": 12.5,
      "num_workers": 32,
      "pin_memory": false
    },
    "train": {
      "num_codebooks": 8,
      "batch_size": 8,
      "grad_accum": 1,
      "max_steps": 110463,
      "warmup_steps": 1100,
      "min_lr_ratio": 0.0,
      "depth_chunk_size": 100,
      "precision": "bfloat16",
      "max_grad_norm": 1.0,
      "seed": 42,
      "scheduler_total_steps": null,
      "per_chip_batch": 2,
      "schedule": "wsd",
      "wsd_anneal_steps": 11000,
      "wsd_anneal_frac": 0.1,
      "weight_decay": 0.01,
      "adam_beta1": 0.9,
      "adam_beta2": 0.999,
      "adam_eps": 1e-08,
      "allow_fresh_optimizer": false,
      "use_scan_layers": true,
      "xla_grad_checkpoint": true,
      "compile_warmup_steps": 1,
      "debug_input_sharding": false,
      "early_stop_patience": 10,
      "enable_clip_grad_norm": true,
      "clip_grad_norm": 1.0,
      "log_grad_norm": true,
      "fsdp_barrier_hook": false,
      "micro_mark_step": true,
      "final_canonical_save": true
    },
    "loss": {
      "text_weight": 0.2,
      "audio_weight": 1.0,
      "text_weight_curriculum_frac": 0.0,
      "composite_text_w": 0.4,
      "composite_audio_w": 0.6,
      "text_padding_weight": 0.01,
      "zero_padding_weight": 0.0,
      "label_smoothing": 0.1,
      "per_codebook_multipliers": [
        1.0,
        1.0,
        1.5,
        1.5,
        2.0,
        2.0,
        2.0,
        2.0
      ],
      "progressive_unmask_fraction": 0.1,
      "unmask_k0": 1
    },
    "optim": {
      "lr_lora": 0.0001716,
      "lr_full_ft": 5e-05,
      "lr_projection": 0.0005,
      "lr_depth": 0.0001,
      "lr_audio_embed": 0.0005,
      "lr_text_embed": 0.0005,
      "lr_depth_blocks": 1e-05
    },
    "logging": {
      "log_every": 25,
      "save_every": 1000,
      "audio_every": 5000,
      "audio_ar_frames": 50,
      "val_every": 250,
      "val_max_batches": 25,
      "val_on_tpu": true,
      "val_per_chip_batch": 8,
      "tpu_telemetry_every": 250,
      "save_dir": "gs://tinyaya-stage2-eu/checkpoints/stage2-v6e16-mh-v03-r2",
      "log_spaced_saves": true,
      "save_at_steps": null,
      "async_checkpoint_upload": true,
      "wandb_project": "tinyaya-stage2-tpu",
      "wandb_run_name": "v0.3-long-horizon-mh-r2",
      "use_wandb": true,
      "wandb_tags": [
        "v0.3",
        "multi-host",
        "wsd"
      ],
      "push_to_hub": true,
      "hub_repo_id": "tiny-aya-translate/tr-hi-s2st-v0.3",
      "hub_private": true,
      "diag_metrics": true,
      "keep_last_n": 0,
      "keep_local_checkpoints": false
    },
    "perf": {
      "enabled": true,
      "warmup_skip_steps": 25,
      "xprof_trace_labels": false
    },
    "lora": {
      "r": 32,
      "alpha": 64,
      "dropout": 0.05,
      "use_rslora": true,
      "target_modules": [
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
        "embed_tokens"
      ],
      "lora_exclude_top": 0,
      "num_full_ft_layers": 0,
      "depth_unfreeze_blocks": 0
    }
  },
  "wandb_run_id": "xzcb60bl",
  "best_val": 3.1171767711639404,
  "patience_left": 10,
  "resumes": 0,
  "files": {
    "rng.pt": 6605,
    "depth_decoder.pt": 1428745845,
    "text_embed.pt": 1090668924,
    "scheduler.pt": 1401,
    "projection.pt": 16778814,
    "optimizer.pt": 768419326,
    "model_audio_embed.pt": 8394415,
    "audio_heads.pt": 8390277,
    "peft_adapter/adapter_config.json": 1122,
    "peft_adapter/README.md": 5202,
    "peft_adapter/adapter_model.safetensors": 1220179520
  }
}