Update Float8Tensor for GRPO training in unsloth #3158

andrewor14 · 2025-10-12T22:21:23Z

Summary: Support a few extra ops called during GRPO loop in unsloth/vllm for Float8Tensor.

Test Plan:

python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_fp8_matmul_lora_variants
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_to_dtype_layout
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_has_compatible_shallow_copy_type
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_transpose

pytorch-bot · 2025-10-12T22:21:28Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3158

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit a437d92 with merge base 6815e57 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

vkuzo

plz clean up _float8_mm_impl

vkuzo · 2025-10-31T11:21:06Z

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

+    input_tensor: Float8Tensor,
+    weight_tensor: Float8Tensor,
+    bias: Optional[torch.Tensor] = None,
+    weight_is_already_transposed: bool = False,


instead of this flag, just transpose at the callsite to match the meaning of matmul

The reason behind this flag is to prevent unnecessary double transpose when we call linear, since the fbgemm op expects the weight to be in the linear format (already transposed). So if we don't have this flag:

1. linear calls _float8_mm_impl(input, weight.t()) 2. _float8_mm_impl calls weight.t() before calling torch.ops.fbgemm.f8f8bf16

If we just transpose the weight for linear, we may end up slowing linear down, is that OK?

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

vkuzo · 2025-10-31T11:28:34Z

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

        )


+def _float8_matmul_impl(


how about

define _float8_mm_impl as the float8 version of torch.mm, as the lowest level shared code of maybe quantizing the input and then choosing a gemm

all other functions (matmul, linear, etc) call _float8_mm_impl

it's a bit confusing to have two different paths for linear and matmul

This was done to avoid double transpose in the linear path (which doesn't happen today, see this comment). I agree that ideally everything should go through _float8_mm_impl, but doing so may add overhead for the linear path, should I go ahead and merge the implementations anyway?

Ok, just did some benchmark on the double transpose thing, seems like it didn't introduce much overhead. I refactored the code to the way you suggested, please have another look

vkuzo · 2025-11-05T01:01:29Z

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

            mm_config = weight_tensor.mm_config
            assert mm_config is not None
            assert not _is_128_128_scaled(weight_tensor), "unimplemented"
+            weight_tensor = weight_tensor.t()


nit: just inline this?

res = torch.ops.fbgemm.f8f8bf16_rowwise( ... wq.t(), ..., w_scale.t(), ... )

Hmm I tried doing this but looks like _is_rowwise_scaled expects a rowwise Float8Tensor, and weight_tensor is colwise scaled at this point (it was previously transposed in linear). Also wq and w_scale are used in the tensorwise case too

how about

if _is_rowwise_scaled(weight_tensor.t()): ...

IMO worth changing since we are inside addmm, and addmm expects the a to be (M, K) and b to be (K, N), changing that is a bit confusing

vkuzo · 2025-11-05T01:03:00Z

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

            scaled_mm_config = weight_tensor.mm_config
            assert scaled_mm_config is not None
-            out_shape = get_out_shape(input_tensor.shape, weight_tensor.shape)
+            out_shape = (*input_tensor.shape[:-1], weight_tensor.shape[1])


technically addmm and mm don't support broadcasting of the gemm operands, would be more correct to enforce x and w to be 2d here and move the 3d shape support to matmul and linear. If you want to do a TODO to punt, sounds ok to me.

vkuzo · 2025-11-05T01:06:52Z

test/quantization/quantize_/workflows/float8/test_float8_tensor.py


        self.assertEqual(sliced_dequantized, sliced_original)

+    def test_to_dtype_layout(self):


why do we need this?

vkuzo · 2025-11-05T01:07:02Z

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

+        self.assertEqual(y_fp8.layout, x_fp8.layout)
+        self.assertEqual(y_fp8.device, torch.device("cpu"))
+
+    def test_has_compatible_shallow_copy_type(self):


why do we need this?

These two are called when the transformer trainer moves the model to a different device before training, e.g. https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/module.py#L938

makes sense, add to test_fp8_matmul_lora_variants? IMO the goal of that test is to test the actual use case e2e, so if the actual use case includes the dtype movement then the test should too.

vkuzo · 2025-11-05T01:07:22Z

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

+        sizes: Tuple,
+    ):
+        _, N, K = sizes
+        self._test_fp8_matmul_model(


is there any need to test the backward?

I think it's OK since the backward just dequantizes and does a basic bf16 matmul? I can still add it if you'd like

**Summary:** Support a few extra ops called during GRPO loop in unsloth/vllm for Float8Tensor. **Test Plan:** ``` python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_fp8_matmul_lora_variants python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_to_dtype_layout python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_has_compatible_shallow_copy_type python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_transpose ``` ghstack-source-id: d806897 Pull Request resolved: #3291

**Summary:** Enable FP8 + RL training using TorchAO for 1.36x faster training and 41% lower memory usage: - We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16 - We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel - For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet (this is in progress: vllm-project/vllm#26327) **Example usage:** ``` model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-8B-Base", max_seq_length = 2048, load_in_4bit = False, fast_inference = True, max_lora_rank = 32, load_in_fp8 = True, # set this to True ) \# the rest is the same as before model = FastLanguageModel.get_peft_model(...) ``` **Initial results:** ``` \# fp8 {'train_runtime': 1752.8985, 'train_samples_per_second': 0.228, 'train_steps_per_second': 0.057, 'train_loss': 0.00016130869633343537, 'epoch': 0.01} \# bf16 {'train_runtime': 2390.0045, 'train_samples_per_second': 0.167, 'train_steps_per_second': 0.042, 'train_loss': 0.00015698353490734007, 'epoch': 0.01} ``` **Requires:** - pytorch/ao#3158 (torchao nightly) - unslothai/unsloth-zoo#351

**Summary:** Enable FP8 + RL training using TorchAO for 1.33x faster training and 40% lower memory usage: - We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16 - We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel - For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet (this is in progress: vllm-project/vllm#26327) **Example usage:** ``` model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-8B-Base", max_seq_length = 2048, load_in_4bit = False, fast_inference = True, max_lora_rank = 32, load_in_fp8 = True, # set this to True ) \# the rest is the same as before model = FastLanguageModel.get_peft_model(...) ``` **Initial results:** ``` \# fp8 {'train_runtime': 1725.4337, 'train_samples_per_second': 0.232, 'train_steps_per_second': 0.058, 'train_loss': 0.00015715716748673002, 'epoch': 0.01} \# bf16 {'train_runtime': 2297.8145, 'train_samples_per_second': 0.174, 'train_steps_per_second': 0.044, 'train_loss': 0.00016081033063528594, 'epoch': 0.01} ``` <img width="1199" height="448" alt="Screenshot 2025-11-11 at 4 10 50 PM" src="https://github.com/user-attachments/assets/b6304afd-89e9-42b1-8064-775807e17b23" /> Test script: https://gist.github.com/andrewor14/5b85119fae46845d07b608d420907423 **Requires:** - pytorch/ao#3158 (torchao nightly) - unslothai/unsloth-zoo#351

**Summary:** Enable FP8 + RL training using TorchAO for 1.33x faster training and 42% less model memory usage: - We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16 - We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel - For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet (this is in progress: vllm-project/vllm#26327) **Example usage:** ``` model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-8B-Base", max_seq_length = 2048, load_in_4bit = False, fast_inference = True, max_lora_rank = 32, load_in_fp8 = True, # set this to True ) \# the rest is the same as before model = FastLanguageModel.get_peft_model(...) ``` **Initial results:** ``` \# fp8 {'train_runtime': 1725.4337, 'train_samples_per_second': 0.232, 'train_steps_per_second': 0.058, 'train_loss': 0.00015715716748673002, 'epoch': 0.01} \# bf16 {'train_runtime': 2297.8145, 'train_samples_per_second': 0.174, 'train_steps_per_second': 0.044, 'train_loss': 0.00016081033063528594, 'epoch': 0.01} ``` <img width="1199" height="448" alt="Screenshot 2025-11-11 at 4 10 50 PM" src="https://github.com/user-attachments/assets/b6304afd-89e9-42b1-8064-775807e17b23" /> Test script: https://gist.github.com/andrewor14/5b85119fae46845d07b608d420907423 **Requires:** - pytorch/ao#3158 (torchao nightly) - unslothai/unsloth-zoo#351

**Summary:** Enable FP8 + RL training using TorchAO for 1.33x faster training and 42% less model memory usage: - We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16 - We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel - For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet (this is in progress: vllm-project/vllm#26327) **Example usage:** ``` model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-8B-Base", max_seq_length = 2048, load_in_4bit = False, fast_inference = True, max_lora_rank = 32, load_in_fp8 = True, # set this to True ) \# the rest is the same as before model = FastLanguageModel.get_peft_model(...) ``` **Initial results:** ``` \# fp8 {'train_runtime': 1725.4337, 'train_samples_per_second': 0.232, 'train_steps_per_second': 0.058, 'train_loss': 0.00015715716748673002, 'epoch': 0.01} \# bf16 {'train_runtime': 2297.8145, 'train_samples_per_second': 0.174, 'train_steps_per_second': 0.044, 'train_loss': 0.00016081033063528594, 'epoch': 0.01} ``` <img width="1199" height="448" alt="Screenshot 2025-11-11 at 4 10 50 PM" src="https://github.com/user-attachments/assets/b6304afd-89e9-42b1-8064-775807e17b23" /> Test script: https://gist.github.com/andrewor14/5b85119fae46845d07b608d420907423 **Requires:** - pytorch/ao#3158 (torchao nightly or 0.15.0+) - unslothai/unsloth-zoo#351

* Enable FP8 + RL training for bf16 models **Summary:** Enable FP8 + RL training using TorchAO for 1.33x faster training and 42% less model memory usage: - We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16 - We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel - For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet (this is in progress: vllm-project/vllm#26327) **Example usage:** ``` model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-8B-Base", max_seq_length = 2048, load_in_4bit = False, fast_inference = True, max_lora_rank = 32, load_in_fp8 = True, # set this to True ) \# the rest is the same as before model = FastLanguageModel.get_peft_model(...) ``` **Initial results:** ``` \# fp8 {'train_runtime': 1725.4337, 'train_samples_per_second': 0.232, 'train_steps_per_second': 0.058, 'train_loss': 0.00015715716748673002, 'epoch': 0.01} \# bf16 {'train_runtime': 2297.8145, 'train_samples_per_second': 0.174, 'train_steps_per_second': 0.044, 'train_loss': 0.00016081033063528594, 'epoch': 0.01} ``` <img width="1199" height="448" alt="Screenshot 2025-11-11 at 4 10 50 PM" src="https://github.com/user-attachments/assets/b6304afd-89e9-42b1-8064-775807e17b23" /> Test script: https://gist.github.com/andrewor14/5b85119fae46845d07b608d420907423 **Requires:** - pytorch/ao#3158 (torchao nightly or 0.15.0+) - unslothai/unsloth-zoo#351 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * _get_inference_mode_context_manager * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py * Update utils.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <[email protected]>

**Summary:** Support a few extra ops called during GRPO loop in unsloth/vllm for Float8Tensor. **Test Plan:** ``` python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_fp8_matmul_lora_variants python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_to_dtype_layout python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_has_compatible_shallow_copy_type python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_transpose ``` ghstack-source-id: d806897 Pull Request resolved: pytorch#3291

* Enable FP8 + RL training for bf16 models (#3440) * Enable FP8 + RL training for bf16 models **Summary:** Enable FP8 + RL training using TorchAO for 1.33x faster training and 42% less model memory usage: - We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16 - We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel - For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet (this is in progress: vllm-project/vllm#26327) **Example usage:** ``` model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-8B-Base", max_seq_length = 2048, load_in_4bit = False, fast_inference = True, max_lora_rank = 32, load_in_fp8 = True, # set this to True ) \# the rest is the same as before model = FastLanguageModel.get_peft_model(...) ``` **Initial results:** ``` \# fp8 {'train_runtime': 1725.4337, 'train_samples_per_second': 0.232, 'train_steps_per_second': 0.058, 'train_loss': 0.00015715716748673002, 'epoch': 0.01} \# bf16 {'train_runtime': 2297.8145, 'train_samples_per_second': 0.174, 'train_steps_per_second': 0.044, 'train_loss': 0.00016081033063528594, 'epoch': 0.01} ``` <img width="1199" height="448" alt="Screenshot 2025-11-11 at 4 10 50 PM" src="https://github.com/user-attachments/assets/b6304afd-89e9-42b1-8064-775807e17b23" /> Test script: https://gist.github.com/andrewor14/5b85119fae46845d07b608d420907423 **Requires:** - pytorch/ao#3158 (torchao nightly or 0.15.0+) - unslothai/unsloth-zoo#351 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * _get_inference_mode_context_manager * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py * Update utils.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <[email protected]> * Update __init__.py * Fix/save torchao model loading logic (#3621) * make loading gpt-oss-BF16 faster. Linked to unsloth-zoo PR #314 * fix model loading and clean merged model directory * revert default quant * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert mapper.py --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Update loader_utils.py * Update loader_utils.py * Add 128x128 PerBlock FP8 + RL (#3629) * Add 128x128 PerBlock FP8 + RL **Summary:** Following #3440, this PR extends torchao FP8 + RL support to also handle 128x128 PerBlock granularity (in addition to PerRow). **Example usage:** ``` model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-8B-Base", max_seq_length = 2048, load_in_4bit = False, fast_inference = True, max_lora_rank = 32, load_in_fp8 = "block", # or "row" or True ) ``` **Initial results:** TBD **Note:** - Requires pytorch/ao#3370 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Version * Update vision.py * Update rl.py * Add torch 2.9.1 * Fix auto installer * Update fp8.py * Float8 * Update fp8.py * Update mapper.py * Update mapper.py * Update loader_utils.py * Update loader.py * Update fp8.py * Versioning * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: andrewor14 <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <[email protected]>

andrewor14 marked this pull request as draft October 12, 2025 22:21

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 12, 2025

andrewor14 added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Oct 12, 2025

andrewor14 force-pushed the unsloth-fp8-rl-test branch from ed3c237 to 19500bf Compare October 12, 2025 22:40

andrewor14 changed the title ~~[HACK] Update Float8Tensor for GRPO training in unsloth~~ [draft] Update Float8Tensor for GRPO training in unsloth Oct 13, 2025

andrewor14 force-pushed the unsloth-fp8-rl-test branch 3 times, most recently from 345bb63 to 9d27057 Compare October 29, 2025 16:32

andrewor14 changed the title ~~[draft] Update Float8Tensor for GRPO training in unsloth~~ Update Float8Tensor for GRPO training in unsloth Oct 29, 2025

andrewor14 requested a review from jerryzh168 October 29, 2025 20:15

andrewor14 marked this pull request as ready for review October 29, 2025 20:15

jerryzh168 reviewed Oct 29, 2025

View reviewed changes

test/quantization/quantize_/workflows/float8/test_float8_tensor.py Outdated Show resolved Hide resolved

jerryzh168 reviewed Oct 29, 2025

View reviewed changes