About Bria.ai
Bria.ai is building a fully copyrighted generative AI platform for developers. AI and engineering teams can use this platform to accelerate visual content creation, increase productivity, and safely develop AI products at scale.
The Challenge
Bria needed to migrate their foundation model training to Amazon SageMaker to create a unified environment for developing and training large-scale AI models (500TB image dataset). They faced significant challenges, including:
- Complex SageMaker onboarding process
- Suboptimal GPU utilization
- Lack of observability and monitoring tools
- Performance optimization needs
- Difficulties with distributed training environments
With that in mind, Bria contacted Automat-it and engaged the experienced MLOps team.
The Solution
Automat-it delivered a comprehensive MLOps solution that included:
Training Infrastructure:
- Unified development and training environment using SageMaker Distributed Data Parallel (DDP)
- Implementation of PyTorch FSDP (Fully Sharded Data Parallel)
- Migration to NVIDIA NeMo Megatron framework
- PyTorch 2.0 integration with optimized Multi-Head Attention
- Custom multi-threaded dataloader development
- Advanced GPU optimization using Parquet and S3 FFM
AWS Marketplace Integration:
- FastAPI frontend with embedded Triton Inference Server
- SageMaker-compliant inference API endpoints
- Multi-GPU instance optimization (g5.12xlarge, p4d.24xlarge)
- PyTorch 2.x optimizations including torch.compile
- Image embedding generation for commercial licensing
The Results
The Automat-it team developed a custom, containerized setup tailored to Bria’s Visual Generative AI models, unlocking more efficient GPU usage and enabling advanced training strategies. Specific results included:
- 57% reduction in training time
- 40% cost reduction
- 11% improvement in inference time
- 3-week SageMaker onboarding – This is dramatically faster than typical implementations
- Near 100% GPU utilization throughout training
Ready to accelerate your AI training and cut costs like Bria.ai? Contact Automat-it today to unlock the full potential of your ML workloads on AWS.