docs(esm2_native_te): add torch.compile troubleshooting note

[melissa-sanabria]

Description

Problem

When running multi-GPU training with torchrun, users may hit the following error:

AssertionError: assert len(self.block_stack) == 1
File "/workspace/bionemo/modeling_esm_te.py", line 259, in forward
with self.get_autocast_context(None, outer=True):

This is a torch.compile incompatibility: TorchDynamo cannot trace through a context manager returned by a method call (self.get_autocast_context(...)). Dynamo has limited support for dynamic context managers, especially ones that wrap torch.autocast in non-trivial ways. As a result, the compilation graph breaks and raises an AssertionError from within the block stack tracker.

Fix

This PR adds a Troubleshooting subsection under ## Commands to Launch Training documenting the known workaround: passing use_torch_compile=false to disable torch.compile at launch time:

torchrun --nproc_per_node=8 train_fsdp2.py use_torch_compile=false

Changes

recipes/esm2_native_te/README.md: added a ### Troubleshooting section with the error signature, root cause explanation, and the workaround command.

Notes

No code changes are included. This is a docs-only PR. The underlying issue — making get_autocast_context compatible with TorchDynamo tracing — would require a separate fix in modeling_esm_te.py

Type of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor
• Documentation update
• Other (please describe):

CI Pipeline Configuration

[ciflow:skip]
Configure CI behavior by applying the relevant labels. By default, only basic unit tests are run.

• ciflow:skip - Skip all CI tests for this PR
• ciflow:notebooks - Run Jupyter notebooks execution tests
• ciflow:slow - Run slow single GPU integration tests marked as @pytest.mark.slow
• ciflow:all - Run all tests, including unit tests, slow tests, notebooks, and every recipe/model directory.

Unit tests marked as @pytest.mark.multi_gpu or @pytest.mark.distributed are not run in the PR pipeline.

For more details, see CONTRIBUTING

Note

By default, only basic unit tests are run. Add appropriate labels to enable an additional test coverage.

Authorizing CI Runs

We use copy-pr-bot to manage authorization of CI
runs on NVIDIA's compute resources.

• If a pull request is opened by a trusted user and contains only trusted changes, the pull request's code will
  automatically be copied to a pull-request/ prefixed branch in the source repository (e.g. pull-request/123)
• If a pull request is opened by an untrusted user or contains untrusted changes, an NVIDIA org member must leave an
  /ok to test comment on the pull request to trigger CI. This will need to be done for each new commit.

Triggering Code Rabbit AI Review

To trigger a code review from code rabbit, comment on a pull request with one of these commands:

• @coderabbitai review - Triggers a standard review
• @coderabbitai full review - Triggers a comprehensive review

See https://docs.coderabbit.ai/reference/review-commands for a full list of commands.

Pre-submit Checklist

• I have tested these changes locally
• I have updated the documentation accordingly
• I have added/updated tests as needed
• All existing tests pass successfully

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3bc34ab2-b496-4950-b00e-6c7136ce18db

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[melissa-sanabria]

Could a maintainer please add the ciflow:skip label? This is a docs-only change.

[pstjohn]

torch compile isn't really supported with TE layers (yet), but thanks for the note!

use_torch_compile should be off by default though:

https://github.com/NVIDIA-BioNeMo/bionemo-recipes/blob/main/recipes/esm2_native_te/hydra_config/defaults.yaml#L11

[melissa-sanabria]

torch compile isn't really supported with TE layers (yet), but thanks for the note!

use_torch_compile should be off by default though:

https://github.com/NVIDIA-BioNeMo/bionemo-recipes/blob/main/recipes/esm2_native_te/hydra_config/defaults.yaml#L11

I tried to run it and I had the error, it worked just after adding the flag. For some reason the default config of use_torch_compile was not taken into account.