How Nota AI Deploys Transformer Models on AWS Inferentia and Trainium

Nota AI  is a company dedicated to the research and development of AI model compression and optimization technologies. Through its proprietary platform, NetsPresso , Nota AI analyzes AI models for compatibility with specific hardware—particularly Neural Processing Units (NPUs)—and modifies the models to align with hardware characteristics. This approach maximizes inference performance and memory efficiency, enabling the deployment of high-performance AI solutions across resource-constrained industries such as mobile, automotive, robotics, and smart cities.

Recently, Nota AI has been expanding its research into lightweight Vision-Language Models (VLMs) and Large Language Models (LLMs), aiming to make domain-specific AI more accessible across diverse applications.

One of the company’s current focal points is cost-efficient LLM inference.

The demand for deploying LLMs in real-world services is rising rapidly. However, the high cost of GPU infrastructure remains a significant barrier. To address this challenge, Nota AI is developing technologies to optimize and reliably deploy LLMs in high-performance NPU environments as a cost-effective alternative to traditional GPUs.

Among these alternatives, AWS Inferentia /Trainium  stand out as promising platforms. More than just high-performance, low-cost chips, they offer dedicated inference infrastructure that is easy to use within the AWS Cloud ecosystem. Notably, the Neuron  SDK simplifies the process of compiling and deploying LLMs, eliminating the need for complex hardware configurations. This significantly reduces both infrastructure costs and operational complexity.

Before diving into full optimization, Nota AI conducted a preliminary study to determine whether popular pre-trained models from the Hugging Face Transformers  library could be ported to AWS Inferentia using the Neuron SDK. The goal was twofold: to establish a solid foundation for future optimization work, and to verify whether these models could run successfully in an NPU environment.

Twelve widely-used Transformer models were selected for testing, including diverse architectures such as GPT, BERT, and OPT, to evaluate the compatibility range of the Neuron SDK.

Using the torch_neuronx.trace() function, Nota AI attempted to port all twelve Transformer-based models to the Neuron SDK. This function converts PyTorch models into a format optimized for AWS Neuron hardware, enabling them to run on AWS Inferentia and Trainium.

The results of the initial tests are as follows:

Out of the twelve models tested, only two were successfully converted, while the remaining models failed during the compilation phase due to RuntimeErrors.

A representative failure case was the meta-llama/Llama-3-2-1B model. While attempting to compile the model using torch_neuronx.trace(), the following error occurred:

Code used:

Error Log:

This error occurs during the tracing process, where torch_neuronx.trace() internally leverages the torch.jit.trace() function to convert a PyTorch model into a static TorchScript graph. torch.jit.trace() traces the execution path of the model using sample inputs and requires the output structure to be statically and consistently defined.

However, models like meta-llama/Llama-3-2-1B return a complex dictionary from their forward() method, mixing multiple data types, such as:

This kind of mixed and nested output structure is not well-suited for tracing-based TorchScript conversion. Since the tracing mechanism cannot reliably predict such dynamic structures at compile time, the compilation fails with a runtime error.

To successfully compile Transformer models like Llama-3 that include diverse output types, the output structure needs to be simplified or customized to return only the necessary elements prior to tracing.

This issue can be easily resolved by setting return_dict=False when loading the model. With this option enabled, the model returns a tuple instead of a dictionary, ensuring consistent output types that are compatible with the tracing process.Using return_dict=False effectively avoids the type inference issues during tracing. For more information on the return_dict option, refer to the Hugging Face documentation .

After applying this change, running torch_neuronx.trace() successfully compiles the model without errors.

To ensure that the converted model produces the same outputs as the original, we tested it using the following sample input:

Output:

The model generated an identical response before and after conversion, confirming that modifying the output structure did not affect the model’s performance. Additionally, it validated that the model could run reliably on NPUs without any degradation in inference quality.

By applying the same approach to the remaining models, we were able to successfully port all 12 Transformer models using the Neuron SDK.

This experiment demonstrated that Transformer models based on Hugging Face can be reliably executed on AWS Inferentia and Trainium. In particular, simple configurations like setting return_dict=False effectively resolved various porting issues. This confirmed that even models with complex output structures can be easily compiled and deployed using the Neuron SDK.

Ultimately, all 12 models were successfully converted, and their output remained consistent with the original models—ensuring no degradation in accuracy or performance after conversion.

Using its AI model optimization platform NetsPresso, Nota AI successfully replaced NVIDIA V100 GPU s with T4 GPU s—doubling the number of channels, improving inference speed (FPS) by 14%, and maintaining accuracy within a 0.1% margin (from 97.14% to 97.03%). These improvements were achieved while reducing monthly operational costs by approximately 85% (from $2,203 to $391).

Building on these quantitative results, Nota AI plans to further enhance memory efficiency and inference speed for a wide range of models on AWS Inferentia and Trainium. This effort aims to empower more organizations to adopt high-performance AI solutions beyond GPUs, while significantly lowering deployment and operating costs.

AWS Trainium and Inferentia offer more than just acceleration for training and inference—they provide a robust and scalable infrastructure for large language models. The Neuron SDK makes it easy to compile and deploy models without complex hardware configurations, making these platforms highly practical for real-world production environments.

Through this experience, we’ve confirmed that AWS’s advanced AI inference platform delivers an outstanding balance of performance, cost-efficiency, and accessibility.

Moving forward, Nota AI will continue optimizing Transformer models and other AI architectures for the AWS Neuron ecosystem, supporting more companies in efficiently operating AI services powered by Inferentia and Trainium.