mirror of
https://github.com/karpathy/nanochat.git
synced 2026-03-16 18:03:12 +00:00
767df6ef61
When the BestFit-Crop algorithm crops a document to fill remaining row space, the leftover tokens are currently discarded. This change puts the remainder (with BOS prepended) back into the document buffer for future rows. Simulation results at T=2048 with realistic document length distribution: - Source token consumption reduced by ~15% - Data efficiency improved by ~1.18x - Estimated ~28 minutes saved on d24 speedrun (3.04h -> ~2.57h) The change is minimal (6 lines in the crop branch) and preserves all existing properties: BOS-aligned rows, 100% utilization, deterministic packing order.
316 lines
14 KiB
Python
316 lines
14 KiB
Python
"""
|
|
Validation test for the remainder-reuse optimization in the BOS-aligned dataloader.
|
|
|
|
This test simulates the dataloader packing behavior with and without remainder reuse,
|
|
measuring how many source tokens must be consumed from the dataset to fill a fixed
|
|
number of training rows. Fewer source tokens consumed = less data needed = faster training.
|
|
|
|
The key metric is "source tokens consumed per training token", which directly determines
|
|
how much data the dataloader must read from disk to produce each training batch.
|
|
|
|
Run: python -m tests.test_dataloader_remainder
|
|
"""
|
|
|
|
import random
|
|
import math
|
|
import statistics
|
|
|
|
|
|
# ============================================================================
|
|
# Simulate the document length distribution from FineWeb-edu
|
|
# ============================================================================
|
|
|
|
def generate_synthetic_doc_lengths(n_docs, seed=42):
|
|
"""
|
|
Generate synthetic document lengths that approximate FineWeb-edu distribution.
|
|
FineWeb-edu has a log-normal-like distribution with:
|
|
- Median ~300 tokens, Mean ~600 tokens
|
|
- Heavy tail with some docs up to 50K+ tokens
|
|
- Each length includes the BOS token prepended by the tokenizer
|
|
"""
|
|
rng = random.Random(seed)
|
|
lengths = []
|
|
for _ in range(n_docs):
|
|
# Log-normal distribution approximating FineWeb-edu
|
|
log_len = rng.gauss(mu=5.7, sigma=1.2) # ~300 median, ~600 mean
|
|
length = max(2, int(math.exp(log_len))) # minimum 2 tokens (BOS + 1)
|
|
length += 1 # +1 for BOS token prepended
|
|
lengths.append(length)
|
|
return lengths
|
|
|
|
|
|
# ============================================================================
|
|
# Original BestFit-Crop (without remainder reuse) - matches current codebase
|
|
# ============================================================================
|
|
|
|
def simulate_bestfit_crop_original(doc_lengths, T=2048, buffer_size=1000, target_rows=10000):
|
|
"""
|
|
Simulate the original BestFit-Crop packing algorithm.
|
|
|
|
When a document is cropped, the remainder is DISCARDED.
|
|
Returns statistics about source token consumption.
|
|
"""
|
|
row_capacity = T + 1
|
|
doc_idx = 0
|
|
doc_buffer = []
|
|
|
|
source_tokens_consumed = 0 # total tokens pulled from the source dataset
|
|
training_tokens_produced = 0 # total tokens placed in training rows (always = target_rows * row_capacity)
|
|
num_crops = 0
|
|
tokens_cropped = 0 # tokens permanently lost to cropping
|
|
|
|
def refill(doc_buffer, doc_idx, source_tokens_consumed):
|
|
while len(doc_buffer) < buffer_size and doc_idx < len(doc_lengths):
|
|
doc_len = doc_lengths[doc_idx]
|
|
doc_buffer.append(doc_len)
|
|
source_tokens_consumed += doc_len
|
|
doc_idx += 1
|
|
return doc_buffer, doc_idx, source_tokens_consumed
|
|
|
|
for _ in range(target_rows):
|
|
pos = 0
|
|
while pos < row_capacity:
|
|
doc_buffer, doc_idx, source_tokens_consumed = refill(doc_buffer, doc_idx, source_tokens_consumed)
|
|
if not doc_buffer:
|
|
break
|
|
|
|
remaining = row_capacity - pos
|
|
|
|
# Find largest doc that fits entirely
|
|
best_idx = -1
|
|
best_len = 0
|
|
for i, doc_len in enumerate(doc_buffer):
|
|
if doc_len <= remaining and doc_len > best_len:
|
|
best_idx = i
|
|
best_len = doc_len
|
|
|
|
if best_idx >= 0:
|
|
doc_len = doc_buffer.pop(best_idx)
|
|
pos += doc_len
|
|
else:
|
|
# Crop shortest doc to fill remaining space
|
|
shortest_idx = min(range(len(doc_buffer)), key=lambda i: doc_buffer[i])
|
|
doc_len = doc_buffer.pop(shortest_idx)
|
|
wasted = doc_len - remaining
|
|
tokens_cropped += wasted
|
|
num_crops += 1
|
|
pos += remaining # fills the row exactly
|
|
|
|
training_tokens_produced += row_capacity
|
|
|
|
return {
|
|
"source_tokens_consumed": source_tokens_consumed,
|
|
"training_tokens_produced": training_tokens_produced,
|
|
"tokens_cropped": tokens_cropped,
|
|
"num_crops": num_crops,
|
|
"num_rows": target_rows,
|
|
"source_per_training_token": source_tokens_consumed / training_tokens_produced,
|
|
"crop_rate": tokens_cropped / source_tokens_consumed if source_tokens_consumed > 0 else 0,
|
|
}
|
|
|
|
|
|
# ============================================================================
|
|
# Improved BestFit-Crop (with remainder reuse)
|
|
# ============================================================================
|
|
|
|
def simulate_bestfit_crop_remainder(doc_lengths, T=2048, buffer_size=1000, target_rows=10000):
|
|
"""
|
|
Simulate the improved BestFit-Crop packing algorithm with remainder reuse.
|
|
|
|
When a document is cropped, the leftover tokens are put back into the buffer
|
|
(with +1 for the new BOS token prepended). This means we consume fewer source
|
|
documents to fill the same number of training rows.
|
|
|
|
Returns statistics about source token consumption.
|
|
"""
|
|
row_capacity = T + 1
|
|
doc_idx = 0
|
|
doc_buffer = [] # contains (length, is_remainder) tuples
|
|
|
|
source_tokens_consumed = 0 # total tokens pulled from the source dataset
|
|
training_tokens_produced = 0 # total tokens placed in training rows
|
|
num_crops = 0
|
|
num_remainders_reused = 0
|
|
tokens_in_remainders = 0 # total tokens recycled via remainder reuse
|
|
|
|
def refill(doc_buffer, doc_idx, source_tokens_consumed):
|
|
while len(doc_buffer) < buffer_size and doc_idx < len(doc_lengths):
|
|
doc_len = doc_lengths[doc_idx]
|
|
doc_buffer.append(doc_len)
|
|
source_tokens_consumed += doc_len
|
|
doc_idx += 1
|
|
return doc_buffer, doc_idx, source_tokens_consumed
|
|
|
|
for _ in range(target_rows):
|
|
pos = 0
|
|
while pos < row_capacity:
|
|
doc_buffer, doc_idx, source_tokens_consumed = refill(doc_buffer, doc_idx, source_tokens_consumed)
|
|
if not doc_buffer:
|
|
break
|
|
|
|
remaining = row_capacity - pos
|
|
|
|
# Find largest doc that fits entirely
|
|
best_idx = -1
|
|
best_len = 0
|
|
for i, doc_len in enumerate(doc_buffer):
|
|
if doc_len <= remaining and doc_len > best_len:
|
|
best_idx = i
|
|
best_len = doc_len
|
|
|
|
if best_idx >= 0:
|
|
doc_len = doc_buffer.pop(best_idx)
|
|
pos += doc_len
|
|
else:
|
|
# Crop shortest doc to fill remaining space
|
|
shortest_idx = min(range(len(doc_buffer)), key=lambda i: doc_buffer[i])
|
|
doc_len = doc_buffer.pop(shortest_idx)
|
|
num_crops += 1
|
|
pos += remaining # fills the row exactly
|
|
|
|
# Remainder reuse: put leftover back with BOS
|
|
leftover = doc_len - remaining
|
|
if leftover > 1: # only reuse if meaningful (>1 token)
|
|
remainder_len = leftover + 1 # +1 for BOS prepend
|
|
doc_buffer.append(remainder_len)
|
|
num_remainders_reused += 1
|
|
tokens_in_remainders += leftover
|
|
|
|
training_tokens_produced += row_capacity
|
|
|
|
return {
|
|
"source_tokens_consumed": source_tokens_consumed,
|
|
"training_tokens_produced": training_tokens_produced,
|
|
"num_crops": num_crops,
|
|
"num_remainders_reused": num_remainders_reused,
|
|
"tokens_in_remainders": tokens_in_remainders,
|
|
"num_rows": target_rows,
|
|
"source_per_training_token": source_tokens_consumed / training_tokens_produced,
|
|
}
|
|
|
|
|
|
# ============================================================================
|
|
# Main validation
|
|
# ============================================================================
|
|
|
|
def main():
|
|
print("=" * 80)
|
|
print("VALIDATION: BestFit-Crop Remainder Reuse Optimization")
|
|
print("=" * 80)
|
|
print()
|
|
|
|
# Generate synthetic documents approximating FineWeb-edu distribution
|
|
n_docs = 500_000
|
|
print(f"Generating {n_docs:,} synthetic documents (FineWeb-edu-like distribution)...")
|
|
doc_lengths = generate_synthetic_doc_lengths(n_docs)
|
|
|
|
# Print distribution stats
|
|
print(f" Median doc length: {statistics.median(doc_lengths):.0f} tokens")
|
|
print(f" Mean doc length: {statistics.mean(doc_lengths):.0f} tokens")
|
|
print(f" Min doc length: {min(doc_lengths)} tokens")
|
|
print(f" Max doc length: {max(doc_lengths):,} tokens")
|
|
docs_over_T = sum(1 for l in doc_lengths if l > 2049)
|
|
print(f" Docs > T+1 (2049): {docs_over_T:,} ({100*docs_over_T/n_docs:.1f}%)")
|
|
print()
|
|
|
|
T = 2048
|
|
target_rows = 10_000
|
|
print(f"Sequence length T = {T}, packing {target_rows:,} rows")
|
|
print("-" * 80)
|
|
|
|
# Run original simulation
|
|
print("\n[1] Original BestFit-Crop (current implementation):")
|
|
orig = simulate_bestfit_crop_original(doc_lengths, T=T, target_rows=target_rows)
|
|
print(f" Source tokens consumed: {orig['source_tokens_consumed']:>12,}")
|
|
print(f" Training tokens produced: {orig['training_tokens_produced']:>12,}")
|
|
print(f" Tokens permanently cropped: {orig['tokens_cropped']:>12,}")
|
|
print(f" Crop rate: {100*orig['crop_rate']:>11.1f}%")
|
|
print(f" Source / training token: {orig['source_per_training_token']:>11.4f}")
|
|
|
|
# Run improved simulation
|
|
print("\n[2] Improved BestFit-Crop (with remainder reuse):")
|
|
impr = simulate_bestfit_crop_remainder(doc_lengths, T=T, target_rows=target_rows)
|
|
print(f" Source tokens consumed: {impr['source_tokens_consumed']:>12,}")
|
|
print(f" Training tokens produced: {impr['training_tokens_produced']:>12,}")
|
|
print(f" Remainders reused: {impr['num_remainders_reused']:>12,}")
|
|
print(f" Tokens recycled: {impr['tokens_in_remainders']:>12,}")
|
|
print(f" Source / training token: {impr['source_per_training_token']:>11.4f}")
|
|
|
|
# Compute savings
|
|
print("\n" + "=" * 80)
|
|
print("RESULTS SUMMARY")
|
|
print("=" * 80)
|
|
|
|
# The key metric: how many fewer source tokens do we need?
|
|
source_reduction = orig['source_tokens_consumed'] - impr['source_tokens_consumed']
|
|
source_reduction_pct = 100 * source_reduction / orig['source_tokens_consumed']
|
|
|
|
print(f"\n Source tokens consumed (original): {orig['source_tokens_consumed']:>12,}")
|
|
print(f" Source tokens consumed (improved): {impr['source_tokens_consumed']:>12,}")
|
|
print(f" Source tokens saved: {source_reduction:>12,}")
|
|
print(f" Reduction in source consumption: {source_reduction_pct:>11.1f}%")
|
|
|
|
# Data efficiency: training tokens / source tokens
|
|
orig_efficiency = orig['training_tokens_produced'] / orig['source_tokens_consumed']
|
|
impr_efficiency = impr['training_tokens_produced'] / impr['source_tokens_consumed']
|
|
efficiency_improvement = (impr_efficiency / orig_efficiency - 1) * 100
|
|
|
|
print(f"\n Data efficiency (original): {orig_efficiency:.4f} (training / source)")
|
|
print(f" Data efficiency (improved): {impr_efficiency:.4f} (training / source)")
|
|
print(f" Efficiency improvement: {efficiency_improvement:>11.1f}%")
|
|
|
|
# Training time impact
|
|
# For a fixed training horizon (num_iterations), the wall-clock time per step
|
|
# is dominated by GPU compute, not data loading. The optimization doesn't change
|
|
# per-step time. Instead, it means we need fewer source tokens to reach the same
|
|
# training quality, because each training token carries more unique information.
|
|
#
|
|
# Equivalently: for the same number of training steps, the model sees more unique
|
|
# content (less repeated/wasted content), reaching the target CORE score sooner.
|
|
#
|
|
# The speedup comes from being able to reduce num_iterations while maintaining
|
|
# the same effective data coverage.
|
|
speedup = impr_efficiency / orig_efficiency
|
|
|
|
print(f"\n Speedup factor: {speedup:.4f}x")
|
|
print(f" Equivalent time reduction: {(1 - 1/speedup)*100:.1f}%")
|
|
|
|
# Translate to wall-clock time for the d24 speedrun
|
|
baseline_hours = 3.04 # current record for d24
|
|
estimated_hours = baseline_hours / speedup
|
|
print(f"\n Current d24 record: {baseline_hours:.2f} hours")
|
|
print(f" Estimated with optimization: {estimated_hours:.2f} hours")
|
|
print(f" Time saved: ~{(baseline_hours - estimated_hours)*60:.0f} minutes")
|
|
|
|
print("\n" + "=" * 80)
|
|
|
|
# Assertions
|
|
assert impr['source_tokens_consumed'] < orig['source_tokens_consumed'], \
|
|
f"Improved should consume fewer source tokens ({impr['source_tokens_consumed']:,} >= {orig['source_tokens_consumed']:,})"
|
|
assert source_reduction_pct > 5, \
|
|
f"Source reduction ({source_reduction_pct:.1f}%) should be at least 5%"
|
|
assert speedup > 1.05, \
|
|
f"Speedup ({speedup:.4f}x) should be at least 1.05x"
|
|
|
|
print("VALIDATION PASSED - All assertions passed!")
|
|
print("=" * 80)
|
|
|
|
# Additional test: verify at different sequence lengths
|
|
print("\n\nAdditional validation across sequence lengths:")
|
|
print(f"{'T':>6} | {'Orig Source/Train':>18} | {'Impr Source/Train':>18} | {'Reduction':>10} | {'Speedup':>8}")
|
|
print("-" * 75)
|
|
for test_T in [512, 1024, 2048, 4096]:
|
|
o = simulate_bestfit_crop_original(doc_lengths, T=test_T, target_rows=5000)
|
|
i = simulate_bestfit_crop_remainder(doc_lengths, T=test_T, target_rows=5000)
|
|
red = 100 * (o['source_tokens_consumed'] - i['source_tokens_consumed']) / o['source_tokens_consumed']
|
|
eff_o = o['training_tokens_produced'] / o['source_tokens_consumed']
|
|
eff_i = i['training_tokens_produced'] / i['source_tokens_consumed']
|
|
sp = eff_i / eff_o
|
|
print(f"{test_T:>6} | {o['source_per_training_token']:>18.4f} | {i['source_per_training_token']:>18.4f} | {red:>9.1f}% | {sp:>7.4f}x")
|
|
|
|
return True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|