Katkat
Back

Optimize LLM Training and Serving

AdobeMachine Learning Engineer · Onsite
HardMachine Learning EngineerOnsiteML System Design

Last updated

This question tests the ability to diagnose hardware utilization bottlenecks in LLM workloads and design architectural optimizations to improve throughput and latency under constrained GPU memory resources.

Scenario

You are designing the training and serving pipeline for a 70B parameter LLM (16-bit precision) on a cluster of H100 GPUs. During initial training runs, your monitoring tools report a Model FLOPs Utilization (MFU) of 35% and a Hardware FLOPs Utilization (HFU) of 62%. You observe significant spikes in HBM (High Bandwidth Memory) usage during the forward pass, and latency for small-batch serving is dominated by memory-bound kernels.

Task

  1. Diagnostic Analysis: Explain the discrepancy between MFU and HFU in this context. What specific hardware or software factors likely contribute to the 35% MFU metric given the HFU performance?
  2. Training Optimization: Propose a strategy to optimize the training throughput. Discuss the trade-offs between activation checkpointing (gradient checkpointing), tensor parallelism, and sequence parallelism when dealing with long context windows (e.g., 32k tokens).
  3. Serving Latency: Describe an architecture for a low-latency serving engine. How would you apply PagedAttention and KV cache quantization (e.g., INT8/FP8) to mitigate the memory bandwidth bottleneck during the autoregressive decoding phase?

Constraints & Assumptions

  • Assume each H100 provides 80GB of VRAM and a peak FP16/BF16 TFLOPS of ~989 (dense).
  • Focus on the impact of the memory wall vs. compute-bound operations.
  • Consider both the communication overhead between nodes and the local compute-to-memory ratio.

Hints

Hint 1

Consider how activation re-computation via checkpointing reduces memory footprint at the expense of extra FLOPs, affecting the MFU calculation.

Hint 2

Analyze the decoding phase of transformer models; note that since it is memory-bandwidth bound, compute utilization (TFLOPS) is often secondary to effective KV-cache management.

Hint 3

Distinguish between HFU (actual FLOPs performed relative to peak) and MFU (theoretical FLOPs required for the model architecture) to identify if your overhead is caused by communication, IO, or re-computation.

Follow-up questions

  • How does the choice of interconnect (e.g., NVLink vs. InfiniBand) specifically impact the scaling efficiency of Tensor Parallelism compared to Pipeline Parallelism?
  • If we implement Speculative Decoding to improve throughput, how does this change the memory access pattern compared to standard autoregressive generation?
  • How would you measure and mitigate 'tail latency' issues caused by asynchronous KV-cache eviction in a multi-tenant serving environment?