fix(utils): propagate non_blocking in TorchAOBaseTensor._to_copy and _get_to_kwargs by Dev-next-gen · Pull Request #4297 · pytorch/ao

Dev-next-gen · 2026-04-19T00:45:47Z

Problem

_get_to_kwargs explicitly discarded the non_blocking argument parsed from torch._C._nn._parse_to, with a comment saying it is "not very useful for most tensor subclasses". As a result, any call to tensor.to(device, non_blocking=True) on a TorchAOBaseTensor subclass silently became a blocking transfer at the inner-tensor level.

This matters in practice for async CPU→GPU offloading workflows such as diffusers enable_group_offload(use_stream=True): the diffusers hook schedules copies with non_blocking=True so that the transfer stream and the compute stream can overlap. Because the flag was dropped, all copies became blocking, negating the overlap benefit.

On AMD ROCm (gfx1xxx) the missing non_blocking also interacts with a separate stream-ordering race (fixed in huggingface/diffusers#13502): the default stream can race ahead of "blocking" copies that the OS scheduler hasn't committed yet, producing device-mismatch errors in the first matmul.

Fix

_get_to_kwargs: include non_blocking in the returned kwargs dict.
TorchAOBaseTensor._to_copy.default: pop non_blocking from kwargs and forward it to every inner .to() call for both tensor_data_names and optional_tensor_data_names.

The change is backward-compatible: when non_blocking=False (the default), behaviour is identical to before.

Tested on

…_get_to_kwargs ## Problem `_get_to_kwargs` explicitly discarded the `non_blocking` argument parsed from `torch._C._nn._parse_to`, with a comment saying it is "not very useful for most tensor subclasses". As a result, any call to `tensor.to(device, non_blocking=True)` on a `TorchAOBaseTensor` subclass silently became a blocking transfer at the inner-tensor level. This matters in practice for async CPU→GPU offloading workflows such as `diffusers` `enable_group_offload(use_stream=True)`: the diffusers hook schedules copies with `non_blocking=True` so that the transfer stream and the compute stream can overlap. Because the flag was dropped, all copies became blocking, negating the overlap benefit. On AMD ROCm (gfx1xxx) the missing non_blocking also interacts with a separate stream-ordering race (fixed in huggingface/diffusers#13502): the default stream can race ahead of "blocking" copies that the OS scheduler hasn't committed yet, producing device-mismatch errors in the first matmul. ## Fix 1. `_get_to_kwargs`: include `non_blocking` in the returned kwargs dict. 2. `TorchAOBaseTensor._to_copy.default`: pop `non_blocking` from kwargs and forward it to every inner `.to()` call for both `tensor_data_names` and `optional_tensor_data_names`. The change is backward-compatible: when `non_blocking=False` (the default), behaviour is identical to before. ## Tested on - 5× AMD RX 7800 XT (gfx1101), ROCm 7.1, PyTorch 2.7 - FLUX.1-dev int8 (`Int8WeightOnlyConfig`) with `enable_group_offload(use_stream=True)` - Companion fix in diffusers: huggingface/diffusers#13502

pytorch-bot · 2026-04-19T00:45:51Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 1 New Failure

As of commit ce00da1 with merge base b3e0db2 ():

NEW FAILURE - The following job has failed:

Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio mslk --index-url https://download.... / linux-job (gh)
RuntimeError: Command docker exec -t 342cef1101ba689dfd5a9782ebc03c540df6399ead13a17c2598b4c9736698a9 /exec failed with exit code 1

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

meta-cla · 2026-04-19T00:45:53Z

Hi @Dev-next-gen!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

meta-cla · 2026-04-19T02:10:22Z

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

meta-cla · 2026-04-19T02:10:26Z

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

Dev-next-gen · 2026-04-28T11:02:14Z

Hi @jerryzh168 @andrewor14 — gentle ping 🙏

PR is ~9+ days old, CLA signed (confirmed by meta-cla bot), but the CI workflows haven't run
yet — they're stuck in "awaiting approval" state for first-time contributors (visible on
https://hud.pytorch.org/pr/pytorch/ao/4297 — all jobs in U/Workflow Startup Failure since open).

Could a maintainer kick off the workflow approval whenever convenient? It's a small fix
(+10/-10 LOC, 1 file) that propagates non_blocking through TorchAOBaseTensor._to_copy /
_get_to_kwargs. Backward-compatible (no-op when non_blocking=False).

Companion PR on diffusers side: huggingface/diffusers#13502 — together they unblock async
CPU↔GPU offloading for FLUX.1-dev int8 on AMD ROCm.

Happy to iterate on review feedback once CI is green. Thanks!

Dev-next-gen · 2026-04-28T23:09:51Z

@pytorchbot label "module: core"

Dev-next-gen · 2026-04-29T19:27:19Z

@jerryzh168 @vkuzo quick note for whoever reviews: the H100 failure (test_cast_to_float8_e4m3fn_saturation_behavior in test/prototype/mx_formats/test_mx_tensor.py:701) appears unrelated to this PR.
The test has an inline # TODO(#1912): make the saturated cast work in eager mode and remove this test — a recent torch nightly likely resolved #1912, breaking this now-obsolete guard. This PR doesn't touch mx_formats.
Happy to file a separate issue if useful.

jerryzh168

looks good, can you add a test for this?

Verifies the contract change in TorchAOBaseTensor._get_to_kwargs: the returned kwargs dict now includes `non_blocking`, propagated from the original `.to(device, non_blocking=...)` call. Covers three cases: explicit True, explicit False, and default (unspecified). Runs on CPU only, no @skip_if_no_cuda needed. Addresses review feedback on PR pytorch#4297.

Dev-next-gen · 2026-04-29T22:01:04Z

@jerryzh168 added the requested test in test/test_utils.py::TestTorchAOBaseTensor::test_get_to_kwargs_non_blocking — covers non_blocking=True, non_blocking=False, and default (unspecified) cases. CPU-only, runs in every matrix slot. CI is back in awaiting workflow approval after the new push — could you re-trigger when you review?

Dev-next-gen · 2026-04-30T16:50:49Z

@jerryzh168 thanks for the quick turnaround on the review! Let me know if anything else is needed from my side to land this.

Dev-next-gen requested review from jerryzh168 and vkuzo as code owners April 19, 2026 00:45

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 Apr 19, 2026

Dev-next-gen mentioned this pull request Apr 28, 2026

fix(group_offloading): synchronize default stream against transfer stream huggingface/diffusers#13502

Open

pytorch-bot Bot added the module: core changes affecting multiple modules, e.g. base config/tensor, observers, quant ops label Apr 28, 2026

jerryzh168 reviewed Apr 29, 2026

View reviewed changes

jerryzh168 approved these changes Apr 29, 2026

View reviewed changes

jerryzh168 merged commit 28e6aca into pytorch:main Apr 30, 2026
18 of 19 checks passed

jerryzh168 added module: not user facing Use this tag if you don't want this PR to show up in release notes and removed module: core changes affecting multiple modules, e.g. base config/tensor, observers, quant ops labels Apr 30, 2026

jainapurva mentioned this pull request Apr 30, 2026

[CI Fix] Propagate dtype in TorchAOBaseTensor._to_copy #4358

Closed

Navigation Menu

Uh oh!

fix(utils): propagate non_blocking in TorchAOBaseTensor._to_copy and _get_to_kwargs - #4297

fix(utils): propagate non_blocking in TorchAOBaseTensor._to_copy and _get_to_kwargs#4297
jerryzh168 merged 2 commits into
pytorch:mainfrom
Dev-next-gen:fix/to-copy-non-blocking-propagation

Dev-next-gen commented Apr 19, 2026

Uh oh!

pytorch-bot Bot commented Apr 19, 2026 •

edited

Loading

Uh oh!

meta-cla Bot commented Apr 19, 2026

Uh oh!

meta-cla Bot commented Apr 19, 2026

Uh oh!

meta-cla Bot commented Apr 19, 2026

Uh oh!

Dev-next-gen commented Apr 28, 2026

Uh oh!

Dev-next-gen commented Apr 28, 2026

Uh oh!

Dev-next-gen commented Apr 29, 2026

Uh oh!

jerryzh168 left a comment

Uh oh!

Dev-next-gen commented Apr 29, 2026

Uh oh!

Dev-next-gen commented Apr 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Dev-next-gen commented Apr 19, 2026

Problem

Fix

Tested on

Related

Uh oh!

pytorch-bot Bot commented Apr 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

❌ 1 New Failure

Uh oh!

meta-cla Bot commented Apr 19, 2026

Action Required

Process

Uh oh!

meta-cla Bot commented Apr 19, 2026

Uh oh!

meta-cla Bot commented Apr 19, 2026

Uh oh!

Dev-next-gen commented Apr 28, 2026

Uh oh!

Dev-next-gen commented Apr 28, 2026

Uh oh!

Dev-next-gen commented Apr 29, 2026

Uh oh!

jerryzh168 left a comment

Choose a reason for hiding this comment

Uh oh!

Dev-next-gen commented Apr 29, 2026

Uh oh!

Dev-next-gen commented Apr 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Apr 19, 2026 •

edited

Loading