NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Increase lru_cache maxsize in YarnRotaryEmbedding to avoid cache misses with deep models

Open
#1,649 2 comments 0 reactions 0 assignees View on GitHub
community-request enhancement waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**
When training deep models (e.g., 40 layers), the current default maxsize=32 for lru_cache in YarnRotaryEmbedding causes persistent cache misses during 1F1B training. Since each layer's YarnRotaryEmbedding instance generates unique hash keys (due to self being part of the cache key), the LRU eviction policy discards early computations when layer count exceeds maxsize.

This becomes particularly impactful when:

Models have ≥32 layers (common in modern architectures)
Hidden sizes are small (CPU-bound scenarios)
Using 1F1B pipeline parallelism
Benchmarks show ~1% end-to-end throughput degradation in 40-layer models due to recomputation.

**Describe the solution you'd like**
Increase default maxsize to accommodate common model depths (suggest 64 or 128)

**Additional context**

​​**Cache Behavior Analysis**​​ (40-layer model, maxsize=32):
Forward pass: Caches layers 1-40 → only 9-40 remain
Backward pass:
Hits 40-9
Recomputes 8-1 (misses) → cache now contains 1-32
Next forward:
Hits 1-32
Recomputes 33-40

**simple case** : simulates 2F&B operations for a 10 layer model with maxsize=8, which will miss cache in last 2 layers

```
import torch.nn as nn
from functools import lru_cache
class MyClass(nn.Module):
@lru_cache(maxsize=8)
def forward(self, n):
print('missed cache')
return n

class LMM(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.layers = [MyClass() for _ in range(10)]

def forward(self):
for i,l in enumerate(self.layers):
print(i)
l(0)
for i,l in enumerate(self.layers[::-1]):
print(9-i)
l(0)

O = LMM()
O.forward()
O.forward()
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.