Training generative imaging models at scale across GPU and AWS Trainium

A Healthcare & Life Sciences business working with Cloud Combinator on AWS. The client is anonymised at their request.

The challenge had three focus areas.

Training data and model artefacts live in Amazon S3. A SageMaker HyperPod cluster provisions the accelerators, whether NVIDIA GPUs or AWS Trainium chips, and the training job runs across them in parallel. On Trainium, the Sana model runs through PyTorch XLA on the AWS Neuron SDK, with gradients synchronised across every worker so each replica stays in step.

The platform is designed to take the client from constrained, single-machine experiments to elastic, multi-accelerator training that scales on demand. Acceptance is defined against successful distributed training of the Sana model on both GPU and Trainium, a documented performance comparison and production-ready pipelines for future runs. The figures below describe the designed scale of the platform and are engagement targets rather than measured production results.

• Elastic training — A SageMaker HyperPod environment that provisions accelerators on demand and recovers long-running jobs automatically, removing the ceiling imposed by local hardware.
• Accelerator choice — The same Sana workload runs on both NVIDIA GPUs and AWS Trainium, giving the client a genuine price and performance choice for every future training run.
• Evidence-based decisions — A documented cost-performance comparison across the two platforms, so scaling and budget decisions rest on measured throughput and efficiency.
• Production-ready pipelines — Containerised training, checkpointing and BF16 optimisation packaged into repeatable pipelines the team can reuse.
• Handover — Runbooks, technical documentation and knowledge-transfer sessions so the client's engineers can operate and extend the platform independently.

With a scalable training foundation in place across GPU and Trainium, the client is positioned to train larger models and iterate faster, no longer limited by local compute and ready to scale whatever it takes on next.