Skip to content

[RL] Add typed Higgs audio rollouts and gated Megatron semantics - #4

Draft
JingwenGu0829 wants to merge 12 commits into
mainfrom
agent/higgs-rl-from-scratch
Draft

[RL] Add typed Higgs audio rollouts and gated Megatron semantics#4
JingwenGu0829 wants to merge 12 commits into
mainfrom
agent/higgs-rl-from-scratch

Conversation

@JingwenGu0829

@JingwenGu0829 JingwenGu0829 commented Jul 12, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add typed, versioned Higgs [L, 8] action streams and decoded-audio artifacts with strict validation from rollout through Megatron training
  • add the zero-shot sglang-omni /generate client, neutral sampling contract, and complete codebook action/logprob/mask parsing
  • add local Whisper and bounded-concurrency SGLang-Omni Qwen3-ASR rewards with WAV validity checks and guaranteed media release
  • add the naive TP1/PP1/CP1/DP1 Megatron Higgs model, strict 399-tensor checkpoint conversion, complete-row teacher forcing, and joint-row GRPO
  • use Megatron's native Adam, scheduler, checkpointing, and weight-update lifecycle; no manual optimizer or LoRA path
  • avoid local Ray GPU reservations for external rollout proxies and clean up custom NCCL update groups on exit

Correctness boundaries

Higgs emits one factorized multi-discrete action per autoregressive row. The trainer gathers float32 selected logprobs for every sampled codebook cell, sums them to one joint row logprob, and applies the GRPO ratio once per row. Forced BOC/EOC cells remain in teacher-forcing inputs but are masked from the objective.

The server logprobs are retained as a measured cross-engine diagnostic. GRPO's old-policy denominator is the mandatory pre-update Megatron recomputation, so an unchanged policy starts at ratio 1.0 despite fixed serving/training kernel drift. The current real-checkpoint diagnostic bound is explicitly configured as --higgs-logprob-parity-atol 3.0; exact parity remains precision follow-up work.

Validation

latest directly affected suite: 210 passed (181 + 29 Ray tests)
pre-commit over every changed/new file: passed
git diff --check: passed
  • loaded the pinned 4,043,480,576-parameter checkpoint and round-tripped all 399 policy tensors
  • saved and restored the native 56.6 GB Megatron model/Adam checkpoint and completed another optimizer step
  • ran two full Miles rollouts on one H200 trainer plus one H200 shared by Higgs and Qwen3-ASR
  • generated and scored eight WAVs per rollout concurrently
  • observed non-constant group rewards and finite native GRPO/Adam steps with gradient norms 23.11 and 23.38
  • passed joint-logprob diagnostics with maximum differences 2.0013 and 2.3616
  • refit the live server after step 0 and observed rollout 1 use unmixed weight version 2
  • produced valid audio after the update; Qwen3-ASR transcribed a fresh version-2 request exactly
  • reran with Ray advertising only the trainer GPU, proving the external proxy is CPU-only, then reused miles-pp_0 from a second driver without restarting Higgs
  • isolated the prior NCCL failure to mixed CUMEM/IPC mode: Miles forced NCCL_CUMEM_ENABLE=0 while external Omni used the default; symmetric direct and Ray matrices pass, and a fresh 399-tensor refit completed over normal P2P/CUMEM with all transport-disable flags unset
  • consume Omni's versioned NCCL refit descriptor and reject backend, NCCL-version, or CUMEM-mode mismatches before communicator rendezvous

Scope

This draft intentionally supports only the pinned Higgs Audio v3 4B checkpoint with one non-sharded Megatron trainer and zero-shot text-to-audio rollouts. Tensor/pipeline/context parallelism, FSDP, LoRA, reference-audio conditioning, exact cross-engine parity, and systematic ASR benchmark batching are deferred.

Depends on sgl-project/sglang-omni#1032.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant