[ryan-li-ryanllli-aws]

Summary

• read the optional TORCHTITAN_COMPILE_BACKEND environment variable during torchtitan.tools.utils initialization
• call torch.compiler.set_default_backend(...) before attention and vision helpers bind backend-less torch.compile calls
• preserve the existing PyTorch default when the setting is absent

Usage

TORCHTITAN_COMPILE_BACKEND=custom_device_backend ./run_train.sh ...

An out-of-tree custom device integration only needs to set this value before launching TorchTitan. It does not need to wrap torch.compile, patch TorchTitan, or install a TorchTitan-specific plugin entry point. Calls that explicitly pass backend= remain unaffected by PyTorch default-backend selection.

Why this is early enough

torchtitan.models.common.attention imports round_up from torchtitan.tools.utils before defining its class-level and module-level compiled helpers. The common vision encoder imports the attention module before defining its own compiled block-mask helper. Configuring the PyTorch default in tools.utils therefore covers all three import-time sites.

Test plan

• Python 3.11 with nightly CPU PyTorch: 13 passed in tests/unit_tests/cpu/tools/test_tools_utils.py
• trailing whitespace, EOF, AST, license, flake8, ufmt, pydoclint, and codespell hooks passed
• git diff --check passed