v9.1.1
Menu

Gradient Accumulation

mediumPROFine-Tuning MechanicsLLM 2: Training & Alignment

Simulate gradient accumulation over multiple micro-batches.

In gradient accumulation, we sum gradients over K micro-batches before performing one optimizer step. The effective gradient is the average: geff=1K∑i=1Kgig_{\text{eff}} = \frac{1}{K} \sum_{i=1}^{K} g_i

Given a parameter vector, micro-batch gradients, and a learning rate, simulate one full optimization step with gradient accumulation.

Input:

  • Line 1: K lr (accumulation steps, learning rate)
  • Line 2: space-separated floats (initial parameter vector)
  • Next K lines: gradient vectors for each micro-batch

Output: Updated parameter vector after one step: θ' = θ - lr * g_eff, rounded to 4 decimal places.

Example

Input
2 0.1
1.0 2.0
0.5 1.0
1.5 3.0
Output
0.9000 1.8000
Reasoning
  • We calculate the effective gradient geffg_{\text{eff}} by summing the micro-batch gradients and dividing by the number of micro-batches KK: geff=12⋅([0.5,1.0]+[1.5,3.0])=12⋅[2.0,4.0]=[1.0,2.0]g_{\text{eff}} = \frac{1}{2} \cdot ([0.5, 1.0] + [1.5, 3.0]) = \frac{1}{2} \cdot [2.0, 4.0] = [1.0, 2.0]
  • Then, we update the parameter vector θ\theta using the formula θ′=θ−lr⋅geff\theta' = \theta - \text{lr} \cdot g_{\text{eff}}: θ′=[1.0,2.0]−0.1⋅[1.0,2.0]=[1.0,2.0]−[0.1,0.2]=[0.9,1.8]\theta' = [1.0, 2.0] - 0.1 \cdot [1.0, 2.0] = [1.0, 2.0] - [0.1, 0.2] = [0.9, 1.8]
  • The final output is the updated parameter vector θ′\theta' rounded to 4 decimal places: [0.9000,1.8000][0.9000, 1.8000]

Constraints

  • 1 <= K <= 10
  • Parameter and gradient vectors have same dimension
  • Round to 4 decimal places

More from LLM 2: Training & Alignment

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Test results

Sign in to run your code against the test cases.