From 05a78b7d5aea72405560cad05c42cf7d580120af Mon Sep 17 00:00:00 2001 From: Banani Ghosh Date: Sat, 7 Nov 2026 16:23:46 -0800 Subject: [PATCH] fix: import torch before spmd_types to avoid circular import On the Neuron SDK 2.33 DLAMI, importing spmd_types before torch triggers the torch_neuronx backend autoload, which loads torch.distributed.fsdp; _fsdp_param_group.py then re-imports the partially initialized spmd_types and fails with "has no attribute 'register_local_autograd_function'". The Qwen3, DeepSeek-V3 and GPT-OSS recipes import spmd_types at module top, so their config_registry fails to import and ConfigManager reports "Config function ... not found" (e.g. the Qwen3-32B command in docs/dense-tutorials.md). Import torch in torchtitan_neuron/__init__.py so it is loaded before any recipe module. Tested on trn3-dev1.48xlarge, DLAMI ami-0dea3e9e3e3a16119 (torch 2.14.0, torch-neuronx 2.14.3.0.1, spmd_types 0.2.1): the four model config_registry modules import, and the dense-tutorials Qwen3-32B eager command runs 10/10 steps (rc=0). --- torchtitan_neuron/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/torchtitan_neuron/__init__.py b/torchtitan_neuron/__init__.py index 3246efdf889..c9a9f5d7a1a 100644 --- a/torchtitan_neuron/__init__.py +++ b/torchtitan_neuron/__init__.py @@ -9,3 +9,7 @@ core torchtitan. Neuron recipes select them by converting the stock model config (for example ``to_nki_moe_config`` for each model's NKI MoE). """ + +# Import torch before spmd_types: importing spmd_types first makes the torch_neuronx +# autoload re-import the partially initialized spmd_types (circular import). +import torch # noqa: F401, E402