Skip to content

[BugFix] Attempting to make tensordict device transfer marking thread safe. - #1701

Closed
coreyjadams wants to merge 5 commits into
pytorch:mainfrom
coreyjadams:main
Closed

[BugFix] Attempting to make tensordict device transfer marking thread safe.#1701
coreyjadams wants to merge 5 commits into
pytorch:mainfrom
coreyjadams:main

Conversation

@coreyjadams

Copy link
Copy Markdown
Contributor

Description

This PR would close #1700 , ideally. But, I am not the maintainer. I do promise, in this day of AI bots all over github, that I'm at least a real person ...

Motivation and Context

The thread-safety issue I raised in #1700 is actually a challenging performance issue in our datapipes in NVIDIA-physicsnemo. Using a hotfix (not this branch, actually, I did it way messier to verify and monkey-patched tensordict ... no judgement please 😅 )I was able to increase our data loading concurrency dramatically from 1 -> 8x or more. Since that's our main bottleneck, it immediately gives a corresponding speedup (up to network bandwidth limits). So, the motivation here is really "I want threads for IO performance, please".

  • I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

  • Bug fix (non-breaking change which fixes an issue)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

  • I have read the CONTRIBUTION guide (required)
  • I have updated the tests accordingly (required for a bug fix or a new feature).

I don't think this needs updates to the docs ... not sure how much you want to advertise thread safety especially if the whole product isn't deliberately thread-safe? Maybe it already is, actually - please, I'm speaking mostly from ignorance!

I also have not yet looked at updating tests, since I do not know how much python threading tests you want to include?

At the end of the day, if you prefer to ignore this PR and go another route, I will not be insulted! I am proposing this as one option that might work to fix #1700, but I leave it up to you to fix or not?

@meta-cla

meta-cla Bot commented May 14, 2026

Copy link
Copy Markdown

Hi @coreyjadams!

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!

@github-actions

Copy link
Copy Markdown
Contributor

PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Attempting to make tensordict device transfer marking thread safe.

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

@coreyjadams coreyjadams changed the title Attempting to make tensordict device transfer marking thread safe. [BugFix] Attempting to make tensordict device transfer marking thread safe. May 14, 2026
@github-actions github-actions Bot added the bug Something isn't working label May 14, 2026
vmoens added 2 commits May 15, 2026 08:46
Regression coverage for pytorch#1700: a unit-level test that ``mark()`` on one
thread does not leak into another, and a stress test that concurrent
``TensorDict(..., device='cpu')`` constructions across 8 threads do not
race on the recorder. The stress test reliably fires the
``RuntimeError("Can only mark one TensorDict at a time.")`` against
pre-fix code in under a second.

Also updates the existing ``test_non_blocking`` assertion, which read
the private ``_has_transfer`` attribute that this PR renamed and moved
onto the per-thread ``_state``.
@github-actions github-actions Bot added the Test label May 15, 2026

@vmoens vmoens left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the fix! LGTM
would approve even if it were a bot!

I added a test inspired by your MRE

Can you sign the CLA? I can't merge unless that flag is green

Will do a minor with it quickly, 0.13 also coming soon

@coreyjadams

Copy link
Copy Markdown
Contributor Author

Awesome! Yes to the CLA - I have to go through the proper channels for NVIDIA <---> Meta CLA, which may take a business day or two. I'll update here when I'm added, I've kicked off the process on the NVIDIA side. We've got a standing CLA I think I just need to have myself added through the right people.

@coreyjadams

Copy link
Copy Markdown
Contributor Author

I think the CLA has gone through on the NVIDIA side. Whenever you'd like, would you try again? Thanks!

@github-actions github-actions Bot added the Compile torch.compile related label May 21, 2026
# Conflicts:
#	tensordict/typedtensordict.py
#	test/test_compile.py
@vmoens

vmoens commented May 21, 2026

Copy link
Copy Markdown
Collaborator
image It doesn't really appear so right now

@vmoens

vmoens commented May 21, 2026

Copy link
Copy Markdown
Collaborator

Superseeded by #1711
thanks @coreyjadams!

@vmoens vmoens closed this May 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working Compile torch.compile related Test

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Tensordict is not Threadsafe for device transfer in all cases

2 participants