In a bid to tackle the challenges of soaring computational expenses during the Rollout phase of large-scale model reinforcement learning, as well as the issues of policy deviation and training collapse that stem from quantization boundary discrepancies when compressing data from BF16 to FP4, Alibaba QianWen, in conjunction with pertinent teams, has unveiled the TRACE method. Rather than separately striving to enhance quantization accuracy for both the training and Rollout stages, this method leverages the actual FP4 quantization results from the Rollout side to steer the rounding direction on the training side. Additionally, it incorporates a streamlined caching mechanism to minimize the overhead associated with information transmission and storage.

Trials conducted on various MoE models reveal that TRACE can reinstate performance to the BF16 Rollout standard when employing a combination of FP4 W/A and FP4 KV Cache Rollout settings. With an output length of 128K, the Rollout decoding throughput can surge by up to 5.4 times, with merely a 7.4% uptick in additional RL step overhead. Moreover, this approach facilitates the model's gradual acclimatization to the FP4 low-precision execution environment during reinforcement learning, culminating in an FP4 policy quality that outperforms schemes involving quantization post-BF16 RL.