job_type: speechlm

# Reconstructed config that reproduces the model behind the public
# stage-2 pretrained BASE checkpoint (exp/opuslm_v2_stage2_pretrain_base,
# step 260000). The original launch script referenced this filename but it
# was never committed; recreated here as the base-tokenizer analog of
# conf/train_stage2_qwen3.yaml (Qwen3-8B-Base backbone/tokenizer).
# Vocab check: 256 special + 151936 (Qwen3-8B-Base) + 8200 (Xcodec 8x1025)
# = 160392 = embed_tokens rows in the checkpoint.

multimodal_io:
    text:
        tokenizer_name: Qwen/Qwen3-8B-Base
    discrete_audio:
        codec_choice: Xcodec
        codec_hf_model_tag: hf-audio/xcodec-hubert-general
        ssl_choice: null
        ssl_hf_model_tag: null
        delay_interleave: true
        stream_weights: [0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125]
    continuous_audio:
        encoder_choice: huggingface
        encoder_hf_model_tag: Qwen/Qwen3-Omni-30B-A3B-Instruct
        attn_implementation: flash_attention_3
        dtype: bfloat16

model:
    model_choice: parallel
    model_hf_tag: Qwen/Qwen3-8B-Base
    model_conf:
        attn_implementation: flash_attention_3
        dtype: bfloat16
        # NOTE: compile_transformer_body / freeze_text_embeddings appear in the
        # original stage configs but are dead keys in the current espnet2/speechlm
        # code (referenced nowhere). transformers 5.x silently drops them; 4.57.x
        # forwards them into the HF constructor and crashes. Omitted here.
    activation_checkpointing: true

preprocessor:
    audio_input: continuous_audio
    audio_output: discrete_audio
    loss_region: assistant
    audio_cfg: 0.05

data_loading:
    batchfy_method: pack
    batch_size: 15000
    save_loader_state: false
    seed: 43
    num_workers: 6

trainer:
    deepspeed_config: conf/deepspeed_stage2.json
    freeze_param: [multimodal_io_dict.discrete_audio, multimodal_io_dict.continuous_audio]
    max_step: 700000
    save_interval: 5000
    log_interval: 1