[Feature] Add free threading support (#1481)

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>

Author: kevmo314

.github/unittest/linux/scripts/setup_env.sh

@@ -40,13 +40,28 @@ eval "$(${conda_dir}/bin/conda shell.bash hook)"
 printf "python: ${PYTHON_VERSION}\n"
 if [ ! -d "${env_dir}" ]; then
     printf "* Creating a test environment\n"
-    conda create --prefix "${env_dir}" -y python="$PYTHON_VERSION"
+    if [ "${PYTHON_VERSION}" == "3.14t" ]; then
+        # Install free-threaded Python 3.14 from conda-forge
+        conda create --prefix "${env_dir}" -y -c conda-forge python-freethreading
+        # Set PYTHON_GIL=0 to keep GIL disabled
+        export PYTHON_GIL=0
+    else
+        conda create --prefix "${env_dir}" -y python="$PYTHON_VERSION"
+    fi
 fi
 conda activate "${env_dir}"
 
+# For free-threaded Python, ensure PYTHON_GIL=0 is set
+if [ "${PYTHON_VERSION}" == "3.14t" ]; then
+    export PYTHON_GIL=0
+fi
+
 # 3. Install Conda dependencies
 printf "* Installing dependencies (except PyTorch)\n"
-echo "  - python=${PYTHON_VERSION}" >> "${this_dir}/environment.yml"
+# Don't add python version constraint for free-threaded builds
+if [ "${PYTHON_VERSION}" != "3.14t" ]; then
+    echo "  - python=${PYTHON_VERSION}" >> "${this_dir}/environment.yml"
+fi
 cat "${this_dir}/environment.yml"
 
 pip install pip --upgrade

.github/workflows/test-linux.yml

@@ -58,7 +58,7 @@ jobs:
   test-cpu:
     strategy:
       matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
+        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.14t"]
       fail-fast: false
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     permissions:

tensordict/_td.py

@@ -2834,8 +2834,6 @@ def _memmap_(
         if inplace:
             self._is_memmap = True
             self._is_shared = False  # since they are mutually exclusive
-            if self._validate_value_cached is not None:
-                delattr(self, "_validate_value_cached")
             self._device = torch.device("cpu")
         else:
             dest._is_memmap = True

tensordict/base.py

@@ -661,7 +661,6 @@ def __getstate__(self) -> dict[str, Any]:
             "_last_op",
             "_cache",
             "__lock_parents_weakrefs",
-            "_validate_value_cached",
         ):
             result.pop(key, None)
         return result
@@ -3503,8 +3502,6 @@ def dtype(self):
         return self._dtype()
 
     def _batch_size_setter(self, new_batch_size: torch.Size) -> None:
-        if self._validate_value_cached is not None:
-            delattr(self, "_validate_value_cached")
         if new_batch_size == self.batch_size:
             return
         if self._lazy:
@@ -5768,17 +5765,13 @@ def clear_device_(self) -> Self:
 
         """
         self._device = None
-        if self._validate_value_cached is not None:
-            delattr(self, "_validate_value_cached")
         for value in self.values():
             if _is_tensor_collection(type(value)):
                 value.clear_device_()
         return self
 
     def _set_device(self, device: torch.device) -> Self:
         self._device = device
-        if self._validate_value_cached is not None:
-            delattr(self, "_validate_value_cached")
         for value in self.values():
             if _is_tensor_collection(type(value)):
                 value._set_device(device=device)
@@ -12964,26 +12957,21 @@ def _validate_key(self, key: NestedKey) -> NestedKey:
             raise KeyError(_GENERIC_NESTED_ERR.format(key))
         return key
 
-    _validate_value_cached: str | None = None
-
     @property
     def _validate_value(self):
         if is_compiling():
             return self._validate_value_generic
-        _validate_value_cached = self._validate_value_cached
-        if _validate_value_cached is None:
-            if self.device:
-                if self.batch_size:
-                    _validate_value_cached = "_validate_value_generic"
-                else:
-                    _validate_value_cached = "_validate_value_batchfree"
+        if self.device:
+            if self.batch_size:
+                method_name = "_validate_value_generic"
             else:
-                if self.batch_size:
-                    _validate_value_cached = "_validate_value_devicefree"
-                else:
-                    _validate_value_cached = "_validate_value_batchfree_devicefree"
-            self._validate_value_cached = _validate_value_cached
-        return getattr(self, _validate_value_cached)
+                method_name = "_validate_value_batchfree"
+        else:
+            if self.batch_size:
+                method_name = "_validate_value_devicefree"
+            else:
+                method_name = "_validate_value_batchfree_devicefree"
+        return getattr(self, method_name)
 
     def _validate_value_generic(
         self,

tensordict/csrc/pybind.cpp

@@ -13,7 +13,7 @@
 
 namespace py = pybind11;
 
-PYBIND11_MODULE(_C, m) {
+PYBIND11_MODULE(_C, m, py::mod_gil_not_used()) {
   m.def("unravel_keys", &unravel_key, py::arg("key")); // for bc compat
   m.def("unravel_key", &unravel_key, py::arg("key"));
   m.def("_unravel_key_to_tuple", &_unravel_key_to_tuple, py::arg("key"));

test/test_tensordict.py

@@ -14127,6 +14127,41 @@ def test_memmap_robust_key_encoding_bijective(self):
             assert legacy == key
 
 
+class TestFreeThreading:
+    """Tests for free-threading (GIL-less Python) compatibility."""
+
+    def test_concurrent_gc_stress(self):
+        """Regression test for free-threading race condition (PR #1481).
+
+        This test exercises concurrent access to TensorDict instances while
+        triggering garbage collection. On Python 3.14t with PYTHON_GIL=0,
+        this would previously cause segfaults due to a race condition in
+        the _validate_value_cached attribute.
+
+        The test passes on all Python versions but only catches the actual
+        race condition on free-threading builds.
+        """
+        import threading
+
+        def gc_stress():
+            for _ in range(100):
+                td = TensorDict(
+                    {"a": torch.randn(5, 5), "b": {"c": torch.randn(5, 3)}},
+                    batch_size=[5],
+                )
+                # Access _validate_value to trigger the code path that had the race
+                _ = td._validate_value
+                td = None
+                gc.collect()
+
+        threads = [threading.Thread(target=gc_stress) for _ in range(8)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+        # If we get here without segfault, the test passes
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)