DES RQS + ACCELERATOR PARALLELISM
JUSTIN ZHANG - SWAN LAB
AGENDA
- 01Book-keeping items
- 02General learning points
- 03Papers + key takeaways
- 04Research directions for IFS DES scheduler
- 05Action items & next papers
BOOK-KEEPING ITEMS
Happy to see the progression with IFS. I'll start recreating their scheduler this week.
Github Lab credentials are still needed to create the Github Organization.
Domain decision for the lab website is still pending.
Sent a request for a VM instance via the school service desk.
itsjira.carleton.ca/servicedeskSEPTEMBER LAB LUNCH
Sending out a message about a September lunch to coordinate a date with everyone.
- Sri Vadlamani
- Michael MacInnis
- Michael Rolbin
- Justin Zhang
- Paul Roode
- Vineela Pulagam
- Kieran Aston
- Noah Kasdorf
- Thank you for the Discord link - I haven't sent a message yet, I'll get to that this week.
- I've created a mock-up of the revamped lab website and would love your thoughts. swan.spawns.dev
VERSIONING - CALVER
Nothing research-relevant on the list today. Since your field is Software Engineering, I wanted to talk about versioning - specifically Calendar Versioning (CalVer).
CALENDAR VERSIONING
CalVer tags releases by date instead of arbitrary semantic numbers.
ACCELERATOR PARALLELISM - TRAINING AT SCALE
Two foundational papers on training large models across multiple accelerators: GPipe (pipeline parallelism) and Megatron-LM (tensor parallelism).
GPIPE
Enables training on multiple accelerators (GPUs, TPUs) in a parallel manner. What makes this paper unique:
MICRO-BATCHES
A batch-splitting algorithm overlaps compute by splitting the forward and backward pass into smaller micro-batches.
ACCUMULATING GRADIENTS
On backpropagation, each weight w is provided a single gradient telling it which direction to push.
Each weight accumulates a collection of gradients across micro-batches:
w_1 + w_2 + w_3 + ... + w_nDrop-in library partitions a model > splits batches into micro-batches > pipelines them across accelerators > accumulates gradients > syncs.
Not relevant - hardware scaling for multi-GPU, and only applicable for training, not inference. Though many papers build on micro-batching for inference.
MEGATRON-LM
Just like last time when we discussed FlashAttention and how matrices are sliced up, Megatron also uses partitioned matrices (block matrix wiki page).
FlashAttention slices a matrix into sub-matrices to fit into L1 caches (SRAM) for IO-aware computation. Megatron slices the matrix so the multiplications can run on different accelerators, then concatenates them.
BLOCK MATRIX MULTIPLICATION
Tiling, block matrix multiplication, and partitioned matrix multiplication all mean the same thing here.
TRAINING LARGER TRANSFORMERS
- The purpose of Megatron-LM is to train larger transformer models across multiple accelerators, pushing beyond single-accelerator memory limits.
- Tailored to the transformer architecture only.
- Achieved via TILING - partitioning each matrix multiplication across multiple accelerators.
Not relevant - also hardware scaling for multi-GPU. Where GPipe revolves around how batches move through the network, Megatron changes the shape of each batch by slicing the matrix multiplications themselves across accelerators.
ADDITIONAL DETAILS
- Both papers train models with high parameter counts and achieve state-of-the-art results.
- I didn't present the benchmarks here as I didn't find the specific results the most interesting part for my research.
- These are both fundamental papers, so a substantial performance boost is expected.
Both contribute to training at scale - GPipe via how batches flow, Megatron via how matrices are shaped. Neither touches inference scheduling, my research direction.
DES TASK ORCHESTRATOR ARCHITECTURE
The orchestrator scores tasks and hands work to workers as they poll.
DIRECTION 1 - MEMORYLESS vs STATEFUL
Right now the orchestrator has “amnesia” - every time a worker polls, it computes the score from scratch and decides which task to hand over.
- Q1Does memoryless scheduling actually treat tenants fairly, or does it drift over time?
- Q2Could we sneak in one tiny counter per tenant (just a number piggybacked on an existing write) and recover most of the benefits of statefulness?
DIRECTION 2 - LEARNING THE WEIGHTS OFFLINE
In the current score formulation, 0.1 and 0.9 are hardcoded weights:
score = 0.1 * (Norm. backlog) + 0.9 * (Norm. wait time)DES already logs every scheduling decision it makes, so we could train an ML model on those historical logs without ever touching production. (This depends on whether they are open to sharing this data - if not, I can synthesize it, but it won't be realistic.)
Can an RL agent look at these logs and find better weights than 0.1 and 0.9?
DIRECTION 3 - PREDICTING TASK DURATION
You mentioned ML task prediction last meeting. Currently “wait time” just means how long a task has been sitting in the queue - but if we could predict how long a task will actually take, we could schedule smarter (e.g. priority to short-but-waited jobs).
- Q1Can a small ML model predict task duration from features like tenant, doc size, task type, etc.?
Note: a small model could add overhead since the score is recomputed on each worker poll - it needs to be really, really fast. - Q2If we swap “raw wait time” for “predicted remaining time”, does fairness get better or worse?
ACTION ITEMS
- 01Send out the form for a lab lunch and coordinate a September date.
- 02Still waiting for GitHub access - I need the credentials to create the GitHub Organization.
- 03Recreate the IFS DES infrastructure - now that I have the spec for their scheduler (Document Enrichment Service), I can start mocking something up.
NEXT PAPERS
- 01GSPMD - mesh parallelism and SPMD partitioning; automatic sharding for distributed computation graphs.
- 02Re-adjust the reading list to be DES relevant.
DISCUSSION
JUSTIN ZHANG - SWAN LAB - MCS MEETING AUG 13