[BugFix] Fix TransformerModule cache invalidation and allocation

[theap06]

Follow-up to #4193, addressing the post-merge review there.

Cache invalidation after weight updates (P1)

Collector.update_policy_weights_ only notified the policy from _maybe_fallback_update, which the legacy weight_updater path bypasses. A compiled TransformerModule skips the parameter-version fallback, so after a VanillaWeightUpdater sync it kept attending to keys and values computed with the previous weights (an updater built over .data misses in eager too). Collector.load_state_dict had the same hole. The notification now happens once in update_policy_weights_, after whichever path applied the weights, and after a checkpoint restore. The multiprocess worker path already routes through this method, so it is covered without further changes.

Cache allocated under inference mode (P2)

A first cached step under torch.inference_mode() allocated inference tensors, and the next ordinary step failed in reset_kv_cache with Inplace update to inference tensor outside InferenceMode is not allowed; is_init=True did not recover. The cache and position counters are now allocated under torch.inference_mode(False), so they accept in-place writes in both modes. Dynamo only traces that context manager from torch 2.2, so on older torch a compiled step keeps allocating in the caller's mode (implement_for dispatch); eager is fixed everywhere.

Backbone contract (P2)

CausalTransformer documented new_kv_cache(batch_size, device=None) while the module always passes dtype=. The contract now documents dtype (the autocast dtype, or None outside autocast; a backbone may ignore it) and the constructor rejects a backbone whose new_kv_cache does not accept the keyword, instead of failing on the first cached step.

Tests

• test_collector_weight_update_restarts_streams now covers the default sync, a legacy VanillaWeightUpdater and load_state_dict, eager and compiled. Against the unfixed code the legacy case fails in both modes and the checkpoint restore fails when compiled.
• test_inference_mode_step_keeps_cache_usable: an inference-mode step followed by ordinary steps and an episode reset matches a module run entirely in ordinary mode.
• test_custom_backbone uses an independent minimal backbone and checks the dtype handed to it outside and under autocast; test_errs covers the constructor rejection.

Labelled ci/olddeps so the torch < 2.2 branch of the allocation runs in CI.

Olddeps

With the label on, the olddeps shard showed the transformer compile tests failing on torch 2.1. That predates this PR (#4193 never ran the suite): the autocast helper's try/except TypeError is not traceable by dynamo, and fullgraph compile of a TensorDictModule needs a newer dynamo than 2.1 (the issubclass in tensordict's dispatch). The second commit dispatches the autocast query with implement_for and skips the compile tests below torch 2.5, in line with the rest of the suite. The remaining red olddeps and GPU jobs are the same ones failing on main at this PR's base.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4432

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 4 Unrelated Failures

As of commit bbcdfdc with merge base 66dc3fc ():

NEW FAILURE - The following job has failed:

• Unit-tests on Linux / tests-cpu (3.13, bulk) / linux-job (gh)
  test/test_configs.py::TestHydraParsing::test_cql_trainer_parsing_with_file

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Unit-tests on Linux / tests-gpu (3.12, 13.0, 2) / linux-job (gh) (trunk failure)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]
• Unit-tests on Linux / tests-olddeps (3.10, 11.8, remainder) / linux-job (gh) (trunk failure)
  test/test_specs.py::TestSpecMasking::test_project[rand_shape2-device0-mult_discrete-shape3]
• Unit-tests on Linux / tests-olddeps (3.10, 11.8, transforms) / linux-job (gh) (trunk failure)
  test/transforms/test_action_transforms.py::TestClosedLoopMultiAction::test_feedback_and_state[group_key1-3]
• Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 2) / linux-job (gh) (trunk failure)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]

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