Katkat
Back

How would you optimize large-scale training/inference?

NVIDIASoftware Engineer · Technical Screen
MediumSoftware EngineerTechnical ScreenML System Design

Last updated

This problem tests your understanding of performance bottlenecks in large-scale machine learning systems, requiring you to prioritize optimization strategies across memory, compute, and communication dimensions.

Scenario

You are designing the serving infrastructure for a transformer-based Large Language Model (LLM) with 175 billion parameters. The model is currently deployed on a cluster of GPUs, but it suffers from high latency and low GPU utilization during peak traffic.

Task

Describe how you would diagnose the performance bottlenecks and identify a tiered optimization strategy to improve throughput and reduce latency. Your design should address:

  1. Memory Management: How to handle the KV cache growth during autoregressive generation.
  2. Computation: Techniques to accelerate the execution of transformer layers.
  3. Communication: Strategies to mitigate overhead in multi-node inference setups.

Requirements

  • Propose at least three distinct optimization techniques (e.g., kernel fusion, quantization, paging, or parallelism schemes).
  • Provide a brief trade-off analysis for each technique regarding accuracy, latency, and memory footprint.
  • Explain how you would profile the system to verify that your proposed changes are effective.

Hints

Hint 1

Consider the difference between memory-bound operations like KV cache lookups and compute-bound operations like matrix multiplications.

Hint 2

Think about how hardware-level optimizations like Tensor Cores or custom CUDA kernels compare to algorithmic changes like KV cache quantization.

Hint 3

Evaluate the impact of scheduling strategies like Continuous Batching versus simple FIFO queuing on GPU utilization.

Follow-up questions

  • How would your optimization strategy change if you were performing fine-tuning instead of inference?
  • If the model exceeds the memory capacity of a single GPU, how would you decide between Pipeline Parallelism and Tensor Parallelism?
  • How does the choice of precision (e.g., FP32 vs. BF16 vs. INT8) affect your throughput and numerical stability?