Arjun Choudhry

Machine Learning Engineer — distributed LLM training, GPU optimization, and scalable ML systems.

Current roleMachine Learning Engineer, Amazon AGI

FocusDistributed LLM training · GPU & framework optimization

Graduate degreeM.S. Computer Science, Virginia Tech

Undergraduate degreeB.Tech, Delhi Technological University

aj11lfc@gmail.com LinkedIn

Professional Experience

Amazon AGI

Jul 2020 — Present

Machine Learning Engineer · AGI · Alexa

Led MFU tuning across Nova models, delivering millions of dollars in compute cost savings — built an autotuner for 5-D parallelism configuration search and fixed a critical PyTorch profiler bug that had blocked performance analysis org-wide.
Led CUDA Graph integration for large-scale model training, resolving critical bugs across both NVIDIA's OSS and AGI-internal codebases to land 10–20% relative throughput improvement across every production model.
Extended NVIDIA's interleaved 1F1B pipeline schedule for MoE models to run on AGI models, achieving 10–15% relative throughput gains at high expert-parallel degree.
Optimized host–device data transfers by deferring them into pipeline idle periods, eliminating compute blocking and improving throughput 5–10% depending on batch and context size.
Implemented stream parallelization of RMSNorm and RoPE for query and key vectors inside Megatron's attention path, achieving 2–5% relative throughput gains.
Built Sharding Dataloader, a distribution-aware loader that materializes only the datasets and fields relevant to a given rank — cutting S3 throttling and improving memory utilization.
Eliminated persistent training stalls with a proactive stall-detection script and a redesign of dataset caching from monolithic to granular, taking MTTR from ~15 minutes to a few seconds.
Designed CPU-based integration testing across 180+ data providers, plus single-host distributed regression testing that mirrors production parallelism configs — catching errors before they reach a fleet run.
Architected and led development of a config-driven, end-to-end model-building framework automating data prep, training, packaging, tuning, and evaluation — enabling single-click runs adopted across multiple teams.
Built the full packaging and tuning automation for the EPA model, delivering a 20% relative reduction in early cut-off rates, and extended the framework to cover the Beam Search Tuner.
Established the team's expertise in the packaging and tuning stack, becoming the reference point for other teams onboarding onto it.

Virginia Tech — Discovery Analytics Center

2018 — 2020

Graduate Research Assistant

Built an end-to-end system that generates natural-language narratives from acyclic causal graphs, letting analysts explore causal structure through prose.
Composed the pipeline from wikification for entity grounding, time-series analysis over node signals, a question-answering engine for targeted follow-ups, and abstractive summarization to compress multi-hop causal chains.
Explored different model architectures as a means of augmenting causality visualizations with textual narratives, on a study wholly funded by DARPA.
Ran the user study evaluating whether textual narratives improve causal comprehension, publishing the findings in IEEE Transactions on Visualization and Computer Graphics (2020).
Prototyped peer-to-peer energy distribution among microgrid units as smart contracts on the EOSIO VM, modeling generation, demand, and settlement between neighboring households.
Took that prototype to first prize at the Design & Build EOSIO Blockchain Hackathon 2019, and integrated it into Virginia Tech's Grand-Prize-winning DOE Solar Decathlon 2019 entry.

MicroStrategy

2019

Hyper-Intelligence Team

Improved search relevance in the Hyper-Intelligence iOS client by combining fuzzy text matching with Bayesian ranking, making partial and misspelled queries resolve to the intended entity.
Built a card recommendation engine using collaborative filtering, surfacing relevant intelligence cards to users based on the behavior of similar accounts.
Implemented geolocation-based entity filtering and map annotation, so cards resolve to the location a user is actually in rather than the full account catalog.
Shipped all three features in Swift inside the production iOS codebase, working within an existing native architecture and its performance constraints.
Partnered with product and design to scope each feature and demo it to the wider Hyper-Intelligence org.

Siemens Healthineers

Jul 2016 — Aug 2018

Senior Software Developer · Bangalore, IN

Built and shipped front-end interfaces for healthcare web applications in AngularJS, JavaScript, jQuery, HTML, and CSS3, serving clinical users in production.
Owned features end to end across the stack — from UI component work through the Node.js services behind them — within a regulated healthcare engineering environment.
Developed a decentralized Ethereum application that democratizes personal health information across patients and healthcare providers, implementing the QoF incentive scheme without a trusted intermediary such as the NHS.
Wrote the Solidity contracts encoding the scoring and payout logic, paired with a Node.js backend bridging existing clinical data into on-chain state.
That project earned "Best Project" in the Impact category at the annual global awards ceremony, and I was recognized with the Solidarity Award (2017) and Best Performer — Campus to Corporate (2016).

PwC

2015

Technology Consultant · Government Reforms & Infrastructure Development (GRID)

Delivered three live client engagements plus one business-development pursuit within the Government Reforms & Infrastructure Development practice.
Advised on energy-efficiency and green-building programs, covering baseline assessment, savings potential, and implementation roadmaps for public-sector clients.
Contributed analysis to standards-and-labeling policy work, translating appliance efficiency data into recommendations for regulatory thresholds.
Produced the client-facing deliverables — models, reports, and presentations — used in stakeholder reviews with government counterparts.

Education

Virginia Tech

2018 — 2020

M.S., Computer ScienceGPA 4.0

Completed the M.S. with a 4.0 GPA, concentrating on machine learning and large-scale data systems.
Coursework spanned Advanced Machine Learning, Graph Machine Learning, Urban Computing, Data Analytics I & II, Full-Stack and Cloud Software Engineering, and Information Security & Blockchain.
Graduate Research Assistant at the Discovery Analytics Center (DAC), researching narrative generation for causal exploration of directed graphs on a DARPA-funded study.
Co-authored four peer-reviewed publications during the program, across causal visualization and blockchain-based microgrid energy systems.
Won first prize at the Design & Build EOSIO Blockchain Hackathon 2019 and contributed to the team's Grand Prize at the DOE Solar Decathlon 2019.

Delhi Technological University (DCE)

2012 — 2016

B.Tech, Electronics Engineering77.17%

Graduated First Class with Distinction (77.17% aggregate) in Electronics Engineering.
Received the Best Paper Presentation Award at ICATETR-2014 for "Effect of Control Techniques and Motor Parameters on a DC Motor Drive under Various Loading Conditions".
Published on time optimization of instruction execution in FPGAs using embedded systems, extending undergraduate coursework into peer-reviewed research.
Served as an organizer for TROIKA 2014 with the IEEE DTU Student Branch, coordinating one of the university's flagship technical festivals.
Won a gold medal at the annual inter-departmental DTU football tournament.

Technical Expertise

Languages

Python

JavaScript

CUDA

C++

Swift

Solidity

Frameworks

PyTorch

MMegatron

TETransformer Engine

NNeMo

Node.js

AngularJS

Concepts

DTDistributed training

LFLLM fine-tuning

MPModel profiling

AASR

Blockchain

GKGPU & kernel optimization

Awards

Grand Prize

2019

DOE International Solar Decathlon

Golden, Colorado

Blockchain energy work plugged into the winning entry.

energy.gov VT News

First Prize

2019

Design & Build EOSIO Blockchain Hackathon

EOSIO

Microgrid energy-distribution prototype on the EOSIO VM.

Solidarity Award

2017

Siemens Healthcare Pvt. Ltd.

Awarded the Solidarity Award for my overall performance and contribution to the team receiving the "Best Project Award" in the Impact Category.

Best Performer

2016

Campus to Corporate (C2C)

Siemens Healthineers

Awarded the best performer in the Campus to Corporate category for the year 2016.

Gold Medal

2015

DTU Football Tournament

Delhi Technological University

Won gold with the team at the annual inter-departmental football tournament.

Best Paper Award

2014

International Conference on Advance Trends in Engineering, Technology and Research

Best Paper Presentation Award at ICATETR-2014 for the paper "Effect of Control Techniques and Motor Parameters on A DC Motor Drive under Various Loading Conditions".

Student of the Year

2008

Times NIE

Times of India

Felicitated with the TOI NIE Student Of The Year Award for the academic session 2007–2008.

Papers

Google Scholar →