Building a Scalable ML Platform with Metaflow for Distributed LLM Training

autodesk · Tech · year n/a · source Agentic: no · promotional bias: low · source-backed · extracted with prompt 0.5.4 (schema v2)

Generated from extraction.json — do not edit by hand. Every claim cites verbatim source evidence; stated means the quote directly asserts the claim, implied means it supports it indirectly.

Problem

Domain: Enterprise ML platform for distributed LLM training

Autodesk needed a unified ML platform that data scientists would actually adopt, replacing bespoke team platforms while supporting reproducible workflows, AWS-managed execution, secure personal IDE instances, production deployment, distributed GPU training, and high-performance storage/networking for large models and datasets.

Constraints:

• Existing bespoke ML platforms created tech debt and adoption friction
• The platform needed strong UX or teams would not migrate
• Autodesk was boxed into AWS-managed services
• Distributed training required multi-node GPU jobs and monitoring
• Ray needed to integrate with a Metaflow-based managed training platform
• Large models and large datasets made inter-node communication and data access bottlenecks

Architecture

Autodesk built AMP around Metaflow and AWS. Users enter through an AM Studio UI backed by SageMaker Studio, SSO, team project domains, and personal Studio instances. Metaflow orchestrates data prep, training, evaluation, versioning, and reproducible experiments, running AWS Batch jobs and production Step Functions. Distributed training uses Metaflow's parallel decorator, AWS Batch multi-node jobs, curated framework examples, GPU profiling, CloudWatch dashboards, and a custom Ray parallel decorator that creates transient Ray clusters. For HPC scale, Autodesk hardened EFA-enabled AMIs and integrated FSx for Lustre with S3 for low-latency training data access.

Diagram confidence: reconstructed (reconstructed from prose).

flowchart LR
  data-scientist[Data Scientist] --> am-studio[AM Studio]
  am-studio[AM Studio] --> sso[SSO]
  am-studio[AM Studio] --> metaflow[Metaflow]
  metaflow[Metaflow] --> aws-batch[AWS Batch]
  metaflow[Metaflow] --> step-functions[Step Functions]
  metaflow[Metaflow] --> metadata-service[Metadata Service]
  metaflow[Metaflow] --> ray-parallel-decorator[Ray Parallel Decorator]
  ray-parallel-decorator[Ray Parallel Decorator] --> ray-cluster[Ray Cluster]
  ray-cluster[Ray Cluster] --> aws-batch[AWS Batch]
  aws-batch[AWS Batch] --> gpu-nodes[GPU Nodes]
  gpu-nodes[GPU Nodes] --> cloudwatch[CloudWatch]
  gpu-nodes[GPU Nodes] --> efa-networking[EFA Networking]
  gpu-nodes[GPU Nodes] --> fsx-for-lustre[FSx for Lustre]
  step-functions[Step Functions] --> jenkins-lambda[Jenkins Lambda]

Components

Data flow

A data scientist logs into AM Studio through SSO, launches a personal SageMaker Studio instance, writes a Metaflow pipeline, runs training/evaluation through AWS Batch, monitors GPU/CPU utilization in Metaflow UI or CloudWatch, and can promote a prod-ready run through GitHub, Jenkins, Lambda, and Step Functions. For Ray jobs, Metaflow creates a transient Ray cluster on AWS Batch, runs the Ray application, and shuts it down after completion.

Why this architecture

Autodesk chose Metaflow because it unified orchestration, compute, versioning, reproducibility, and AWS Batch integration while offering a user experience strong enough to move teams off bespoke platforms. AWS Batch handled infrastructure, Ray handled distributed hardware usage, EFA mitigated inter-node bottlenecks, and FSx for Lustre avoided repeated S3 reads for terabyte-scale training data.

Models used: 3 billion parameter model · 6 billion GPT-J model · DeepSpeed · Ray Train · PyTorch Lightning · Hugging Face Accelerate · TensorFlow Distributed

Scale

• The early training infrastructure had rolled out to 50 users.
  “um training infrastructure we recently just rolled out to 50 users so we quickly had to put together just a pretty basic UI for now but it does the job um basically it authenticates and authorizes um the data scientists um and”transcript, stated
• Ray Train fine-tuning used 16 A10 GPU nodes and completed in about 50 minutes.
  “dpeed and Ray Trin um the fine tuning job completed in about 50 minutes um just using 16 A10 um GPU nodes and the ray logs are displayed in the metap UI so users can monitor their Ray application seamlessly there are some caveats though um with this integration”transcript, stated
• The Ray parallel decorator was used to fine-tune a six billion parameter model.
  “the ray parallel decorator and used it to fine-tune a six billion rounder model as well as other variety of other applications and use cases um using dpeed and Ray Trin um the fine tuning job completed in about 50 minutes um”transcript, stated
• EFA provides a path to hundreds or thousands of nodes with low-latency communication.
  “pattern or a way a way forward to scale up to hundreds or thousands of nodes with low latency inter node communication um so some of the work that we did here was security Harden the Deep learning Ami the aw has that comes with the EA driver installed and then we”transcript, stated
• An A100 GPU can attach up to four EFA network devices.
  “familiar with that um and then with an a100 GPU you can attach up to four EFA network devices um so given that users will be able to leverage EFA um by running their training jobs with um the with NCC and then we've also open- sourced the integration with EFA in”transcript, stated
• FSx for Lustre provides hundreds of gigabytes per second throughput for terabytes of training data.
  “at with throughput levels of hundreds or even gigabyte hundreds of gigabytes per second um so FSX for luster will be um really useful uh when working with um our terabytes of um training data so that instead of reading it in directly from S3 and saving and saving it to”transcript, stated

  “our terabytes of um training data so that instead of reading it in directly from S3 and saving and saving it to memory or dis space um we can instead use FSX for lusters to sync with S3 and”transcript, stated

Alternatives considered (roads not taken)

Let teams keep bespoke ML platforms

Rejected because: Existing tech debt meant the new platform needed strong UX to motivate migration.

Treat Ray as competing with Metaflow

Rejected because: The team realized Ray was complementary: Batch handled infrastructure and Ray optimized hardware usage.

Read terabytes of training data directly from S3 into memory or disk

Rejected because: FSx for Lustre gives parallel low-latency access synced with S3.

Patterns

durable-execution (workflow)

Metaflow versions flow runs, data snapshots, artifacts, and experiments, then compiles prod flows into Step Functions.

batch-offline-processing (training)

AWS Batch multi-node jobs run distributed training through Metaflow parallel decorators.

observability-tracing (operations)

Metaflow UI, GPU profiling, CloudWatch dashboards, and node logs expose distributed training progress and utilization.

latency-optimization (infrastructure)

EFA and FSx for Lustre reduce inter-node communication and training data access bottlenecks.

Tradeoffs

Choose Metaflow as the orchestration backbone

Chose Metaflow over Other orchestration tools and bespoke platforms.

Rationale: Metaflow covered data, compute, orchestration, versioning, AWS Batch integration, reproducibility, and experiment tracking with a user-friendly workflow.

Integrate Ray through Metaflow rather than replacing Metaflow

Chose Ray parallel decorator on Metaflow/AWS Batch over Separate Ray launcher path.

Rationale: This kept Ray available to users without fragmenting the managed training platform.

Use FSx for Lustre with S3

Chose Parallel filesystem synced with S3 over Direct S3 reads into memory or disk.

Rationale: FSx improved low-latency, simultaneous multi-node access for terabytes of training data.

Gotchas (what bit them)

A technically strong platform can fail if users do not adopt it.

Root cause: Some teams already had bespoke ML platforms, so switching required strong UX.

Resolution: Autodesk centered the platform around AM Studio, personal Studio instances, and familiar Metaflow workflows.

Ray integration was initially a blocker for a research team.

Root cause: A team had invested in Ray while the platform had already chosen Metaflow.

Resolution: Autodesk built a Ray parallel decorator so Metaflow could orchestrate transient Ray clusters on AWS Batch.

Larger models made inter-node communication a bottleneck.

Root cause: Distributed training nodes must exchange data and synchronize model updates.

Resolution: The team security-hardened EFA-enabled deep learning AMIs and integrated them with AWS Batch.

Evaluation

Autodesk validates flows through unit tests, optional flow runs, prod-ready tags, metadata checks, GPU/CPU utilization monitoring, CloudWatch dashboards, and distributed training benchmarks.

Metrics: GPU utilization · CPU utilization · node logs · flow run metadata tags · training time for 6B model · number of GPU nodes

Outcomes

• AMP rolled out to 50 users in its early phase
• The Ray parallel decorator fine-tuned a six billion parameter model in about 50 minutes on 16 A10 GPU nodes
• Autodesk integrated EFA and FSx for Lustre for higher-scale distributed training

The transcript frames the platform as early, with 50 users, but shows a roadmap toward hundreds or thousands of nodes and terabytes of training data.