perf(megatron): keep policy logits in model precision - #2764

Merged

yueming-yuan merged 4 commits into

radixark:mainfrom

YolandaLyj:fix/chunked-bf16-logits

Aug 31, 2026

Commits on Aug 26, 2026

perf(megatron): keep policy logits in model precision
YolandaLyj
committed
c82de0a

Commits on Aug 31, 2026

fix(loss): apply temperature scaling in fp32 inside each logits chunk
yueming-yuan
committed
1b96101
test(loss): add chunked policy snapshot config
yueming-yuan
committed
94d53ec
fix(loss): admit fp16 logits in the chunked log-prob path
yueming-yuan
committed
1be6f92