[RL] Add typed Higgs audio rollouts and gated Megatron semantics - #4
Draft
JingwenGu0829 wants to merge 12 commits into
Draft
[RL] Add typed Higgs audio rollouts and gated Megatron semantics#4JingwenGu0829 wants to merge 12 commits into
JingwenGu0829 wants to merge 12 commits into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
[L, 8]action streams and decoded-audio artifacts with strict validation from rollout through Megatron training/generateclient, neutral sampling contract, and complete codebook action/logprob/mask parsingCorrectness 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.0despite 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
23.11and23.382.0013and2.36162miles-pp_0from a second driver without restarting HiggsNCCL_CUMEM_ENABLE=0while external Omni used the default; symmetric direct and Ray matrices pass, and a fresh 399-tensor refit completed over normalP2P/CUMEMwith all transport-disable flags unsetScope
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.