PyTorch code runs on Trainium with little or no change, as Amazon validates its new backend on 841 representative models

PyTorch developers can now run existing code on AWS Trainium without rewriting their models. Such code was previously written for other accelerators, and moving it to Trainium required extra adaptation work. Amazon states that PyTorch code already running on other accelerators runs on Trainium with minimal or no code changes.

TorchNeuron adds an opt-in adaptive eager mode alongside standard PyTorch's two native execution modes, eager and compiled. Eager mode executes operations immediately in program order and is useful for debugging, while compiled mode is efficient but pays compile time first. A Python interpreter on the CPU accumulates operator-dispatch instructions into small graphs, which TorchNeuron compiles, caches and queues, with multiple graph collections and compilations happening in parallel.

An adaptive algorithm selects subgraphs according to graph and tensor size so that compile time overlaps with execution of already compiled snippets. Larger graphs on larger tensors take longer to compile, so the right subgraphs must be chosen to hide that cost. Amazon says this overlap effectively hides all compile time for a great variety of models while keeping accurate source correlation.

Compiled NEFF files are cached at several levels and compilation is shared across accelerators. Converting operations to NEFFs is expensive, and accelerators running identical parts of a sharded model all need the same NEFFs. In tests with a large language model spread across accelerators, these optimizations cut the time to the first training step by about a third.

The scope is Amazon's release description of TorchNeuron: an open-source backend that translates PyTorch code into operations executable on Trainium, with adaptive eager mode, multi-level NEFF caching and a shared compilation service; the only quantified result given is a roughly one-third cut in time to the first training step in a test with a large language model spread across accelerators.

Amazon describes adaptive eager as approaching fully compiled performance and hiding compile time for a great variety of models, which is an approximation rather than equality; the roughly one-third cut in time to first training step comes from a single large language model test spread across accelerators, with no statement about other model scales.