File size: 5,041 Bytes
d911efa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Same global sample stream across ranks, without dropped tail targets."""
import hashlib
import json
import numpy as np
def benchmark_phases(world):
assert world in (8, 16)
return [{'name': 'strong', 'offset': 0, 'global_batch': 256, 'updates': 40, 'warmup': 10},
{'name': 'weak', 'offset': 10240, 'global_batch': world * 32, 'updates': 12, 'warmup': 4, 'replicated_base_world': 8}]
def step_indices(phase, step, rank, world, samples):
if phase.get('replicated_base_world'):
# Benchmark only: duplicate the same eight real rank streams on the
# second pair of nodes, giving exactly equal per-GPU input/token load.
base = phase['replicated_base_world']
local_count = phase['global_batch'] // world
start = phase['offset'] + step * base * local_count
return list(range(start + rank % base, start + base * local_count, base))
start = phase['offset'] + step * phase['global_batch']
end = min(samples, start + phase['global_batch'])
if start >= samples: return []
return list(range(start + rank, end, world))
def microbatches(examples, max_padded_tokens=12000, max_examples=16):
"""Length sort only within one already fixed optimizer batch."""
examples = sorted(examples, key=lambda ex: len(ex['input_ids']))
batches, current = [], []
for row in examples:
n = len(row['input_ids'])
if current and (n * (len(current) + 1) > max_padded_tokens or len(current) >= max_examples):
batches.append(current); current = []
current.append(row)
if current: batches.append(current)
return batches
def dynamic_weights(local_frames, local_samples, global_frames, global_samples, world):
"""Preserve one global per-channel token mean across microbatches/ranks.
The inherited loss normalizes by sum(weights), so the caller multiplies
its return value by returned scale. DDP subsequently averages gradients.
This is algebraically the ordinary inherited objective on the global batch.
"""
weights = [world * (local_frames + local_samples) / (global_frames + global_samples)]
weights += [(32. / 12) * world * local_frames / global_frames] * 12
return weights, sum(weights) / 33.
def check_plan(samples=16384):
unions = {}
for world in (8, 16):
phases = benchmark_phases(world)
for phase in phases:
seen = []
for step in range(phase['updates']):
global_ids = sorted(i for rank in range(world) for i in step_indices(phase, step, rank, world, samples))
if phase['name'] == 'weak':
start = phase['offset'] + step * 256
expect = sorted(list(range(start, start + 256)) * (world // 8))
else:
expect = list(range(phase['offset'] + step * phase['global_batch'], phase['offset'] + (step + 1) * phase['global_batch']))
assert global_ids == expect
seen.extend(global_ids)
assert len(seen) == len(set(seen)) * (world // 8 if phase['name'] == 'weak' else 1)
unions[(world, phase['name'])] = seen
assert unions[(8, 'strong')] == unions[(16, 'strong')]
for step in range(12):
for rank in range(16):
assert step_indices(benchmark_phases(16)[1], step, rank, 16, samples) == step_indices(benchmark_phases(8)[1], step, rank % 8, 8, samples)
# Test mean/gradient coefficients independently, including unequal lengths.
rng = np.random.default_rng(71)
channels = rng.normal(size=(263, 13))
frame_groups = [13, 49, 81, 120]
sample_groups = [1, 3, 5, 7]
expected = channels.mean(axis=0)
actual = np.zeros(13)
start = 0
for frames, sample_count in zip(frame_groups, sample_groups):
weights, scale = dynamic_weights(frames, sample_count, sum(frame_groups), sum(sample_groups), 4)
# Audio-channel linear coefficients recover the exact global mean after DDP.
audio_coefficient = weights[1] / sum(weights) * scale / 4
actual[1:] += channels[start:start + frames, 1:].mean(axis=0) * audio_coefficient
start += frames
assert np.allclose(actual[1:], expected[1:] * (32 / 12) / 33)
return {'status': 'PASS', 'strong_unique_samples': len(unions[(8, 'strong')]),
'strong_union_sha256': hashlib.sha256(np.asarray(unions[(8, 'strong')], dtype='<i8').tobytes()).hexdigest(),
'worlds': [8, 16], 'global_batch_strong': 256, 'per_gpu_weak': 32,
'strong_warmup_updates': 10, 'strong_steady_updates': 30, 'weak_warmup_updates': 4, 'weak_steady_updates': 8,
'weak_input_rule': 'Exact same eight real per-GPU streams, replicated once across the additional eight GPUs at four nodes. Benchmark-only repetitions, explicitly counted as presentations.',
'loss_weighting': 'Exact global per-channel supervised-token mean, independent of rank and microbatch partition.'}
if __name__ == '__main__': print(json.dumps(check_plan(), indent=2))
|