| { |
| "step": 31000, |
| "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.016290783882141, |
| "patience_left": 9, |
| "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 |
| } |
| } |
