Три конфигурации torchtitan для Qwen3 на одном trn2.48xlarge

TorchNeuron делает Trainium обычным устройством neuron в PyTorch, и torchtitan на нём ведёт себя как на GPU: претрейн с нуля или дообучение опубликованных весов. В walkthrough AWS на одном trn2.48xlarge сравнивают три конфига Qwen3 при общем запуске torchrun --nproc_per_node=64 , по одному процессу на каждое из 64 логических NeuronCore при LNC=2. Три сценария, одна команда Transformers v5 , модель через Hub и AutoModelForCausalLM , чтение и запись safetensors без лишней обвязки. Нативный Qwen3 , в конфиге меняют model_spec на model_registry("4B", attn_backend="sdpa") вместо реализации из Hub. MoE с NKI , рецепты torchtitan-neuron с mixture-of-experts и override registry, чтобы подставить своё ядро под Trainium. Для воспроизведения в примере зафиксированы torch 2.14.0, torch_neuronx 2.14.3, transformers 5.16.0 и torchtitan-neuron release-2.33. Веса Qwen/Qwen3-4B-Instruct-2507 кладут в локальную папку и передают в конфиг. В eager Neuron компилирует операции по ходу; с torch.compile , уже собранный граф. Чекпоинты Hugging Face Первый сценарий стартует с transformers_modeling_backend_full() : ID модели в конфиге, датасет c4_test , 10 шагов обучения и batch 1 на ранг. Для дообучения включают checkpoint, initial_load_path на каталог с весами и флаги initial_load_in_hf / last_save_in_hf , на выходе снова стандартные safetensors. На первом шаге loss для опубликованных весов Qwen3-4B-Instruct , 6,54, со случайной инициализации , 12,43. После экспорта AutoModelForCausalLM подхватывает state dict без лишних ключей, кроме ожидаемого lm_head.weight , в этой модели он связан с эмбеддингами токенов. Нативная сборка и урезанный MoE Во втором варианте берут заготовку qwen3_32b() и подменяют спецификацию на 4B , в той сборке torchtitan не было отдельного билдера qwen3_4b() . Глобальный batch в тесте , 64 при шардировании data parallel на все ранги. Третий путь , урезанный Qwen3-30B-A3B: 4 слоя, 64 эксперта, top_k=4 , tensor parallel = 1, без expert parallelism , чтобы влезть на один инстанс. Там же , регистрация своего rotary-ядра из nkilib через wrap_nki и autograd, с перестановкой layout тензоров под ожидания Hugging Face. Источник: Train Qwen3 on AWS Trainium with torchtitan .