FP8 training

[JosephChenHub]

does trl trainer support FP8 training?

https://huggingface.co/docs/accelerate/en/usage_guides/low_precision_training

[shirinyamani]

Hi @JosephChenHub , Now we are working to bring support for fp8 in GRPO trainer when generating using vllm server mode. here: #3428

I understand what you are mentioning is sth more broad as bringing the support thro accelerate so that it can be used with all the trainers, however, the case based on my understanding can be different especially for online trainers like GRPO, PPO.

We will forsure see what is the best with least complication that we can bring the support.

[JosephChenHub]

Hi @JosephChenHub , Now we are working to bring support for fp8 in GRPO trainer when generating using vllm server mode. here: #3428

I understand what you are mentioning is sth more broad as bringing the support thro accelerate so that it can be used with all the trainers, however, the case based on my understanding can be different especially for online trainers like GRPO, PPO.

We will forsure see what is the best with least complication that we can bring the support.

thanks for your implementation. but how about the weight synchronization when using FP8? I found that the training dtype cannot be bfloat16 when using FP8 inference

[mccatec]

Hi @JosephChenHub , Now we are working to bring support for fp8 in GRPO trainer when generating using vllm server mode. here: #3428

I understand what you are mentioning is sth more broad as bringing the support thro accelerate so that it can be used with all the trainers, however, the case based on my understanding can be different especially for online trainers like GRPO, PPO.

We will forsure see what is the best with least complication that we can bring the support.

You mean sth like what Nvidia has done in their llama-nemotron work?

"We implement custom vLLM weight loaders capable of loading BF16 weights supplied by the training stage, and casting to FP8 weights and scaling factors at runtime. Because vLLM does not support directly initializing models in FP8, we also implement meta-weight tensor initialization to avoid materializing the full BF16 inference engine, which would cause an out-of-memory error in the GPU"

[JosephChenHub]

Hi @JosephChenHub , Now we are working to bring support for fp8 in GRPO trainer when generating using vllm server mode. here: #3428
I understand what you are mentioning is sth more broad as bringing the support thro accelerate so that it can be used with all the trainers, however, the case based on my understanding can be different especially for online trainers like GRPO, PPO.
We will forsure see what is the best with least complication that we can bring the support.

You mean sth like what Nvidia has done in their llama-nemotron work?

"We implement custom vLLM weight loaders capable of loading BF16 weights supplied by the training stage, and casting to FP8 weights and scaling factors at runtime. Because vLLM does not support directly initializing models in FP8, we also implement meta-weight tensor initialization to avoid materializing the full BF16 inference engine, which would cause an out-of-memory error in the GPU"

exactly

[shirinyamani]

Hi @JosephChenHub , Now we are working to bring support for fp8 in GRPO trainer when generating using vllm server mode. here: #3428
I understand what you are mentioning is sth more broad as bringing the support thro accelerate so that it can be used with all the trainers, however, the case based on my understanding can be different especially for online trainers like GRPO, PPO.
We will forsure see what is the best with least complication that we can bring the support.

You mean sth like what Nvidia has done in their llama-nemotron work?

"We implement custom vLLM weight loaders capable of loading BF16 weights supplied by the training stage, and casting to FP8 weights and scaling factors at runtime. Because vLLM does not support directly initializing models in FP8, we also implement meta-weight tensor initialization to avoid materializing the full BF16 inference engine, which would cause an out-of-memory error in the GPU"

Hi @mccatec and @JosephChenHub
Thanks for your reply. I'll explain the problem to the best of my knowledge.
So currently in grpo which is an online trainer, we are using vllm for generation. The way its implemented as explained in out vllm_integration doc, is sending the generation request to vllm_server via vllm_client which is the trainer.
The good news is that vllm naturally supports --quantization which can be set to fp8, meaning for generation/inference purposes we can use fp8. However the bad news is that the weights which are being sent from the client_trainer to the vllm_server for completion generation are not in fp8!!

To overcome this conflict, one very easy solution is perhaps converting the weights loaded from trainer to the precision compatible with the --quantization setup (e.g. fp8), this perhaps should be somewhere in this class.

so technically, the goal here is to have the support for fp8 for inference so that overall memory usage be lower!

[Krish-Goyani]

Hi @JosephChenHub , Now we are working to bring support for fp8 in GRPO trainer when generating using vllm server mode. here: #3428

I understand what you are mentioning is sth more broad as bringing the support thro accelerate so that it can be used with all the trainers, however, the case based on my understanding can be different especially for online trainers like GRPO, PPO.

We will forsure see what is the best with least complication that we can bring the support.

Are there any updates regarding this?

Does the GRPO trainer now support the FP8 training??

[ehartford]

why you guys talking about vllm?

I wanna train, not inference.

does trl yet support fp8 for training?

[romitjain]

Hi @shirinyamani
Just checking if there has been any update on this?

[albertvillanova]

Hey @JosephChenHub, sorry for the long silence here — we're catching up on the backlog.

Two separate things got tangled in this thread, so to be clear:

1. FP8 training itself isn't really a TRL feature. TRL trainers subclass transformers.Trainer, which delegates mixed precision to Accelerate — so you enable it exactly as in the doc you linked, via accelerate launch --mixed_precision fp8 (with a TransformerEngine / MS-AMP / torchao backend). Nothing TRL-specific is needed for plain SFT/DPO/etc. @ehartford — this is the answer to your "I want to train, not inference" question too.

2. The genuinely open piece is FP8 generation for the online trainers (GRPO/PPO) with weight sync to vLLM — casting the BF16 trained weights to FP8 at load time, Nemotron-style. The first attempt (🛠️ quantization support for vllm generation #3428) was closed unmerged, so this is still up for grabs.

Keeping this open as the tracking thread; 👍 / +1 to signal demand for (2).

[ehartford]

Thank you for the education! 😊

[albertvillanova]

A small update on where item (2) actually stands, for anyone who wants to pick it up. The blocker isn't exposing a --quantization fp8 flag on trl vllm-serve: it's the weight sync. update_named_param broadcasts the raw training-dtype tensor and calls load_weights straight into the running engine (

trl/trl/scripts/vllm_serve.py

Lines 117 to 149 in 922dc58

	def update_named_param(self, name: str, dtype: str, shape: Sequence[int]) -> None:
	"""
	Receives updated weights from the client process and updates the named parameter in the model.
	Args:
	name (`str`):
	Name of the weight tensor being updated.
	dtype (`str`):
	Data type of the weight tensor as a string (e.g., `"torch.float32"`).
	shape (`Sequence[int]`):
	Shape of the weight tensor.
	"""
	import torch
	from transformers import is_torch_xpu_available
	if self.communicator is None:
	raise RuntimeError("Communicator not initialized. Call `init_communicator` first.")
	dtype = getattr(torch, dtype.split(".")[-1])
	# Allocate memory for the incoming weight tensor on the correct device.
	weight = torch.empty(shape, dtype=dtype, device=self.device)
	if is_torch_xpu_available():
	# Use XCCL to broadcast the updated weights from the client (src) to all workers.
	self.communicator.broadcast(weight, root=self.client_rank)
	self.communicator.barrier()
	else:
	# Use NCCL to broadcast the updated weights from the client (src) to all workers.
	self.communicator.broadcast(weight, src=self.client_rank)
	self.communicator.group.barrier()
	# Load the received weights into the model.
	self.model_runner.model.load_weights(weights=[(name, weight)])

), so an FP8-initialized engine would receive BF16 tensors with no accompanying scale recomputation — exactly the custom-loader problem @mccatec quoted from Llama-Nemotron, and why both #3428 and #5367 stalled. In colocate mode quantization is auto-derived today and is only ever "bitsandbytes" (QLoRA) or None, with no user-facing knob (

trl/trl/generation/vllm_generation.py

Lines 341 to 348 in 922dc58

	quantization = None
	if is_bitsandbytes_available():
	for _, module in model.named_modules():
	if isinstance(module, bnb.nn.Linear4bit):
	quantization = "bitsandbytes"
	break
	elif isinstance(module, bnb.nn.Linear8bitLt):
	raise ValueError("vLLM does not support in-flight 8-bit quantization.")

). If your goal is mainly lower memory, note that trl vllm-serve --kv_cache_dtype fp8 already works today and needs no weight-sync changes (

trl/trl/scripts/vllm_serve.py

Lines 287 to 292 in 922dc58

	kv_cache_dtype: str = field(
	default="auto",
	metadata={
	"help": "Data type to use for KV cache. If set to 'auto', the dtype will default to the model data type."
	},
	)

).

Keeping this open as the tracking thread; 👍 here to signal demand for the full FP8-weights path.

[yashb98]

Hi @albertvillanova, I'd like to take this on — the weight-sync casting issue in update_named_param/load_weights you described. Will dig into the FP8 scale-recomputation piece and follow up here with a plan / questions once I've had a closer look.