Deterministic replay notebook

Python starts when you choose Run.

100 identical loss logs, 100 different model states

Aryan Putta · November 6, 2026

A Megatron-inspired CPU replay lab shows why matching seeds and rounded loss curves cannot establish bitwise reproducibility.

The experiment uses CPU float32 regression and Adam with synthetic data. Run the cells in order; results and figures are written to the working directory. Dependencies: NumPy and Matplotlib.

Method: CPU float32 regression with Adam on synthetic data. The controlled comparisons change gradient-reduction order and checkpoint state to test exact replay.

Sources:

https://developer.nvidia.com/blog/scale-bitwise-deterministic-pretraining-with-nvidia-megatron-core/
https://docs.pytorch.org/docs/stable/notes/randomness.html
https://docs.nvidia.com/cuda/floating-point/index.html

import os
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
os.environ.setdefault('OMP_NUM_THREADS', '1')
"""A deterministic CPU training replay lab inspired by NVIDIA Megatron Core.
Run: OPENBLAS_NUM_THREADS=1 python experiment.py
Dependencies: Python, NumPy, Matplotlib. No downloads, network, model or GPU.
"""
import os
os.environ.setdefault('MPLCONFIGDIR','/tmp/engineering-blog-mpl')
import copy,json,hashlib,csv
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
ROOT=Path(__file__).resolve().parent if '__file__' in globals() else Path.cwd()
STEPS=100;CUT=40;BATCH=64;D=24
COLORS={'Same seed + full state':'#eeeeee','Full checkpoint':'#eeeeee','Reversed reduction':'#333333','Optimizer reset':'#9a9a9a','Sampler RNG reset':'#536b7a'}

def digest(a):
    a=np.ascontiguousarray(a)
    return hashlib.sha256(str(a.shape).encode()+str(a.dtype).encode()+a.tobytes()).hexdigest()

def make_data():
    r=np.random.default_rng(20261106)
    x=r.normal(size=(512,D)).astype(np.float32)
    truth=r.normal(size=D).astype(np.float32)
    y=(x@truth+.2*np.sin(3*x[:,0])+.03*r.normal(size=512)).astype(np.float32)
    return x,y

def init_state():
    return {'w':np.random.default_rng(17).normal(0,.01,size=D).astype(np.float32),'m':np.zeros(D,np.float32),'v':np.zeros(D,np.float32),'step':0,'rng':copy.deepcopy(np.random.default_rng(99).bit_generator.state)}

def state_digest(s):
    return hashlib.sha256((''.join(digest(s[k]) for k in ['w','m','v'])+json.dumps({'step':s['step'],'rng':s['rng']},sort_keys=True)).encode()).hexdigest()

def reduce_rows(x,reverse=False):
    out=np.zeros(x.shape[1],np.float32)
    for row in (x[::-1] if reverse else x):
        out=np.add(out,row,dtype=np.float32)
    return out

def step(state,data,reverse=False):
    x,y=data
    rng=np.random.default_rng();rng.bit_generator.state=copy.deepcopy(state['rng'])
    idx=rng.choice(len(x),BATCH,replace=False)
    state['rng']=copy.deepcopy(rng.bit_generator.state)
    xb=x[idx];yb=y[idx]
    pred=np.sum(xb*state['w'],axis=1,dtype=np.float32)
    residual=pred-yb
    per_example=(np.float32(2/BATCH)*residual[:,None]*xb).astype(np.float32)
    g=reduce_rows(per_example,reverse)
    input_hash=digest(per_example);grad_hash=digest(g)
    state['step']+=1;t=state['step']
    b1=np.float32(.9);b2=np.float32(.999);lr=np.float32(.04)
    state['m']=(b1*state['m']+(np.float32(1)-b1)*g).astype(np.float32)
    state['v']=(b2*state['v']+(np.float32(1)-b2)*g*g).astype(np.float32)
    mh=state['m']/np.float32(1-float(b1)**t)
    vh=state['v']/np.float32(1-float(b2)**t)
    state['w']=(state['w']-lr*mh/(np.sqrt(vh)+np.float32(1e-8))).astype(np.float32)
    # Float64 evaluation avoids hiding loss changes inside its own float32 reduction.
    residual_eval=x.astype(np.float64)@state['w'].astype(np.float64)-y
    loss=float(np.mean(residual_eval**2))
    return {'step':t,'loss':loss,'printed_loss':f'{loss:.4f}','weights':state['w'].copy(),'weight_hash':digest(state['w']),'state_hash':state_digest(state),'gradient_input_hash':input_hash,'gradient_hash':grad_hash,'sample_hash':digest(idx)}

def run(data,state=None,steps=STEPS,reverse=False):
    s=init_state() if state is None else copy.deepcopy(state);trace=[]
    for _ in range(steps):trace.append(step(s,data,reverse))
    return s,trace

def serial_trace(trace):
    return [{k:v for k,v in t.items() if k!='weights'} for t in trace]

def run_experiment():
    data=make_data();final,ref=run(data)
    _,same=run(data)
    checkpoint,prefix=run(data,steps=CUT)
    restored,tail=run(data,checkpoint,STEPS-CUT)
    _,reverse=run(data,reverse=True)
    broken_optimizer=copy.deepcopy(checkpoint);broken_optimizer['m'][:]=0;broken_optimizer['v'][:]=0
    # Model the common partial-restart behavior: reset both Adam moments and its step counter.
    broken_optimizer['step']=0
    _,badopt=run(data,broken_optimizer,STEPS-CUT)
    # Report training positions on the common x-axis; the optimizer counter remains reset internally.
    for j,t in enumerate(badopt):t['step']=CUT+j+1
    broken_rng=copy.deepcopy(checkpoint);broken_rng['rng']=copy.deepcopy(np.random.default_rng(99).bit_generator.state)
    _,badrng=run(data,broken_rng,STEPS-CUT)
    traces={'Reference':ref,'Same seed + full state':same,'Full checkpoint':prefix+tail,'Reversed reduction':reverse,'Optimizer reset':prefix+badopt,'Sampler RNG reset':prefix+badrng}
    rows=[]
    for name,trace in traces.items():
        if name=='Reference':continue
        for a,b in zip(ref,trace):
            rows.append({'scenario':name,'step':a['step'],'exact_weights_match':a['weight_hash']==b['weight_hash'],'exact_state_match':a['state_hash']==b['state_hash'],'same_4dp_loss':a['printed_loss']==b['printed_loss'],'max_weight_error':float(np.max(np.abs(a['weights'].astype(float)-b['weights']))),'loss':b['loss'],'reference_loss':a['loss']})
    return data,traces,rows,checkpoint,final,restored

def validate(data,traces,rows,checkpoint,final,restored):
    tests=[]
    def check(name,ok):assert ok,name;tests.append(name)
    check('same-state repeat matches at every step',all(r['exact_state_match'] for r in rows if r['scenario']=='Same seed + full state'))
    check('full checkpoint matches at every step',all(r['exact_state_match'] for r in rows if r['scenario']=='Full checkpoint'))
    check('full resumed final arrays are byte equal',all(final[k].tobytes()==restored[k].tobytes() for k in ['w','m','v']))
    check('full resumed RNG and step equal',final['step']==restored['step'] and final['rng']==restored['rng'])
    a=traces['Reference'][0];b=traces['Reversed reduction'][0]
    check('reduction comparison has identical ordered inputs',a['gradient_input_hash']==b['gradient_input_hash'])
    check('reduction order changes first gradient bytes',a['gradient_hash']!=b['gradient_hash'])
    check('changed order is hidden by rounded logging somewhere',any(not r['exact_weights_match'] and r['same_4dp_loss'] for r in rows if r['scenario']=='Reversed reduction'))
    for name in ['Optimizer reset','Sampler RNG reset']:
        check(name+' matches prefix',all(r['exact_state_match'] for r in rows if r['scenario']==name and r['step']<=CUT))
        check(name+' breaks replay after restore',any(not r['exact_state_match'] for r in rows if r['scenario']==name and r['step']>CUT))
    check('training loss decreases',traces['Reference'][-1]['loss']<traces['Reference'][0]['loss'])
    check('fingerprint detects permutation',digest(np.array([1,2,3]))!=digest(np.array([3,2,1])))
    # Different dtypes and shapes must not accidentally share a byte-only digest.
    check('fingerprint includes shape and dtype',digest(np.array([1,2],np.int32))!=digest(np.array([[1,2]],np.int32)))
    return tests

def summarize(rows,traces):
    out=[]
    for name in COLORS:
        r=[r for r in rows if r['scenario']==name]
        bad=[v['step'] for v in r if not v['exact_weights_match']]
        out.append({'scenario':name,'first_weight_mismatch':min(bad) if bad else None,'weight_mismatch_steps':len(bad),'hidden_weight_mismatch_steps':sum(not v['exact_weights_match'] and v['same_4dp_loss'] for v in r),'maximum_weight_difference':max(v['max_weight_error'] for v in r),'final_loss':r[-1]['loss']})
    return out

def make_figures(rows,traces,out=ROOT):
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':11,'axes.spines.top':False,'axes.spines.right':False,'axes.edgecolor':'#555555','axes.labelcolor':'#333333','text.color':'#222222','xtick.color':'#555555','ytick.color':'#555555','axes.titleweight':'normal','axes.titlesize':12,'svg.fonttype':'none','figure.facecolor':'white','axes.facecolor':'white','grid.color':'#dddddd','grid.linewidth':.6})
    fig,axs=plt.subplots(2,1,figsize=(9.2,6.4),gridspec_kw={'height_ratios':[1.1,1.5]})
    names=list(COLORS);matrix=np.array([[int(r['exact_weights_match']) for r in rows if r['scenario']==n] for n in names])
    from matplotlib.colors import ListedColormap
    from matplotlib.patches import Patch
    axs[0].imshow(matrix,aspect='auto',extent=(.5,100.5,4.5,-.5),cmap=ListedColormap(['#536b7a','#eeeeee']),vmin=0,vmax=1,interpolation='nearest')
    axs[0].set_yticks(range(5),names);axs[0].set_xticks([1,20,40,60,80,100]);axs[0].tick_params(axis='y',length=0);axs[0].set_title('Exact weight agreement with the reference',loc='left',pad=30);axs[0].axvline(CUT+.5,color='#333333',lw=.8)
    axs[0].legend(handles=[Patch(facecolor='#eeeeee',label='Matching bytes'),Patch(facecolor='#536b7a',label='Different bytes')],loc='lower left',bbox_to_anchor=(0,1.01),ncol=2,frameon=False,fontsize=9,borderaxespad=0)
    styles={'Reversed reduction':('#333333','-'),'Optimizer reset':('#9a9a9a','--'),'Sampler RNG reset':('#536b7a',':')}
    for n,(c,ls) in styles.items():
        aa=[r for r in rows if r['scenario']==n and r['max_weight_error']>0]
        axs[1].plot([r['step'] for r in aa],[r['max_weight_error'] for r in aa],color=c,ls=ls,label=n,lw=1.6)
    axs[1].set_yscale('log');axs[1].set_xlabel('Training step');axs[1].set_ylabel('Maximum weight difference');axs[1].axvline(CUT+.5,color='#bbbbbb',ls=':',lw=1);axs[1].grid(axis='y');axs[1].legend(frameon=False,fontsize=9,ncol=3,loc='upper center',bbox_to_anchor=(.5,-.22))
    fig.text(.03,.012,'Float32 CPU study · checkpoint after step 40 · zero differences omitted from the logarithmic plot',fontsize=9,color='#666666')
    fig.tight_layout(rect=(0,.055,1,1),h_pad=2.5)
    for name in ['replay','cover']:
        fig.savefig(out/(name+'.svg'),bbox_inches='tight');fig.savefig(out/(name+'.png'),dpi=180,bbox_inches='tight')
    plt.close(fig)
    ref=traces['Reference'];rev=traces['Reversed reduction']
    fig,ax=plt.subplots(figsize=(7,4))
    ax.plot(range(1,STEPS+1),[x['loss'] for x in ref],color='#aaaaaa',lw=2.2,label='Reference');ax.plot(range(1,STEPS+1),[x['loss'] for x in rev],color='#333333',lw=1.2,ls='--',label='Reversed reduction')
    ax.set_xlabel('Training step');ax.set_ylabel('Evaluation mean squared error');ax.legend(frameon=False,fontsize=10);ax.grid(axis='y');fig.tight_layout();fig.savefig(out/'hidden-drift.svg',bbox_inches='tight');fig.savefig(out/'hidden-drift.png',dpi=180,bbox_inches='tight');plt.close(fig)

def main():
    data,traces,rows,checkpoint,final,restored=run_experiment();tests=validate(data,traces,rows,checkpoint,final,restored);summary=summarize(rows,traces)
    (ROOT/'results.json').write_text(json.dumps({'config':{'steps':STEPS,'checkpoint_step':CUT,'batch':BATCH,'dimensions':D,'rows':512,'data_seed':20261106,'weights_seed':17,'sampler_seed':99,'dtype':'float32','optimizer':'Adam','learning_rate':.04},'summary':summary,'tests':tests,'traces':{k:serial_trace(v) for k,v in traces.items()}},indent=2))
    with (ROOT/'comparisons.csv').open('w',newline='') as f:w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    np.savez(ROOT/'checkpoint.npz',w=checkpoint['w'],m=checkpoint['m'],v=checkpoint['v'])
    (ROOT/'checkpoint-state.json').write_text(json.dumps({'step':checkpoint['step'],'rng':checkpoint['rng'],'data_sha256':digest(data[0])+':'+digest(data[1]),'note':'Load checkpoint.npz with allow_pickle=False. Recreate the documented fixed data and configuration.'},indent=2))
    # Verify the file round-trip independently of the in-memory replay.
    loaded=np.load(ROOT/'checkpoint.npz',allow_pickle=False);meta=json.loads((ROOT/'checkpoint-state.json').read_text());s={k:loaded[k] for k in ['w','m','v']};s.update(step=meta['step'],rng=meta['rng']);disk_final,_=run(data,s,STEPS-CUT)
    assert state_digest(disk_final)==state_digest(final);tests.append('saved checkpoint file round-trip matches reference')
    make_figures(rows,traces)
    saved=json.loads((ROOT/'results.json').read_text());saved['tests']=tests;(ROOT/'results.json').write_text(json.dumps(saved,indent=2))
    (ROOT/'checks.json').write_text(json.dumps({'passed':tests,'count':len(tests)},indent=2))
    print(json.dumps({'checks_passed':len(tests),'summary':summary,'reference_final_loss':traces['Reference'][-1]['loss']},indent=2))

Run the study and its assertions

The printed results are computed, not illustrative. Change the constants in the previous cell to explore a different controlled setup.

{
  "checks_passed": 15,
  "summary": [
    {
      "scenario": "Same seed + full state",
      "first_weight_mismatch": null,
      "weight_mismatch_steps": 0,
      "hidden_weight_mismatch_steps": 0,
      "maximum_weight_difference": 0.0,
      "final_loss": 0.020268479118598137
    },
    {
      "scenario": "Full checkpoint",
      "first_weight_mismatch": null,
      "weight_mismatch_steps": 0,
      "hidden_weight_mismatch_steps": 0,
      "maximum_weight_difference": 0.0,
      "final_loss": 0.020268479118598137
    },
    {
      "scenario": "Reversed reduction",
      "first_weight_mismatch": 1,
      "weight_mismatch_steps": 100,
      "hidden_weight_mismatch_steps": 100,
      "maximum_weight_difference": 1.1920928955078125e-07,
      "final_loss": 0.020268477811632834
    },
    {
      "scenario": "Optimizer reset",
      "first_weight_mismatch": 41,
      "weight_mismatch_steps": 60,
      "hidden_weight_mismatch_steps": 1,
      "maximum_weight_difference": 0.1605449914932251,
      "final_loss": 0.024593437931813907
    },
    {
      "scenario": "Sampler RNG reset",
      "first_weight_mismatch": 41,
      "weight_mismatch_steps": 60,
      "hidden_weight_mismatch_steps": 1,
      "maximum_weight_difference": 0.02873385325074196,
      "final_loss": 0.02016254975299807
    }
  ],
  "reference_final_loss": 0.020268479118598137
}

Inspect the generated figures

These are rendered from the results just computed.

Loss and parameter differences between two float32 training runs with changed reduction order.

Model-state divergence after checkpoint restoration variants.

Python runs locally in this page. Saved results and figures are visible before execution.

Download full experiment notebook