{"step": 41000, "restored_from_hub": true, "source_repo": "fushinguyenex/IMF", "checkpoint": "step_0041000.pt", "trainable_parameters": 14992182, "note": "Metrics before this step remain in the previous Modal workspace"}
{"resume_from_step": 41000}
{"step": 41001, "lr": 1.77396025983391e-05, "window_steps": 1, "elapsed_seconds": 25.122098667000003, "steps_per_second": 0.039805591612996266, "eta_seconds": 226073.76590433303, "images_seen": 41985024, "images_per_second": 40.76092581170818, "effective_batch_size": 1024, "loss": 0.34101641178131104, "edge_loss": 0.1725691556930542, "diag_loss": 0.16844725608825684, "jvp_rms": 0.76953125, "grad_norm": 0.14997200667858124, "grad_norm_postclip": 0.14997200667858124, "clip_fraction": 0.0, "ema_updates": 41001, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.34101641178131104, "updates_nfe_3": 1, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.895012855529785, "gradient_groups": {"t_emb": {"weight_norm": 16.476406916050635, "grad_norm": 0.016168814749080263, "grad_rms": 4.708200924785152e-05, "grad_to_weight": 0.0009813313564943139, "update_to_weight": 0.00015766120394309556, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.240280296274523, "grad_norm": 0.013997224757716357, "grad_rms": 4.075855130472968e-05, "grad_to_weight": 0.0008618832004351109, "update_to_weight": 0.00011253604207629143, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.057381762387267, "grad_norm": 0.0034498390355066186, "grad_rms": 6.176202615920352e-05, "grad_to_weight": 0.0006821393356467062, "update_to_weight": 6.275091777852822e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.830638123577607, "grad_norm": 0.01528453657317228, "grad_rms": 4.89105796397333e-05, "grad_to_weight": 0.0011912530324649547, "update_to_weight": 0.0001413048977955577, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.061590607521455, "grad_norm": 0.04032836463277654, "grad_rms": 0.0009175020644546593, "grad_to_weight": 0.00796752794918837, "update_to_weight": 4.8407424306601647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63334824759529, "grad_norm": 0.07129177650093073, "grad_rms": 0.000171787911987034, "grad_to_weight": 0.004286074904446102, "update_to_weight": 0.00015689318456363897, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.552468928169933, "grad_norm": 0.024436467142409563, "grad_rms": 3.371377912123787e-05, "grad_to_weight": 0.0009952753514891012, "update_to_weight": 0.00016219035859577446, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.819434200096254, "grad_norm": 0.01276580788295766, "grad_rms": 1.5497393373012116e-05, "grad_to_weight": 0.00041421292162845065, "update_to_weight": 0.0001361478908286072, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.58817246985865, "grad_norm": 0.008839003004888865, "grad_rms": 7.93479277538175e-06, "grad_to_weight": 0.0002555498707713304, "update_to_weight": 0.00016035676654794785, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01082638271199, "grad_norm": 0.00745571504297393, "grad_rms": 5.997092484907456e-06, "grad_to_weight": 0.00020144686762402434, "update_to_weight": 0.0001733819149955561, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.655452894276248, "grad_norm": 0.003039499937344634, "grad_rms": 5.491438299593651e-06, "grad_to_weight": 0.00012327901460087363, "update_to_weight": 0.00011410106013083344, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.25576061790646, "grad_norm": 0.007700976734720823, "grad_rms": 6.194371624457871e-06, "grad_to_weight": 0.00020670566395628644, "update_to_weight": 0.0001683364476197874, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.353409860998205, "grad_norm": 0.019536428956998135, "grad_rms": 8.673497556942531e-06, "grad_to_weight": 0.000308372177596477, "update_to_weight": 0.00017971743270335145, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.73741858363756, "grad_norm": 0.028534908584078696, "grad_rms": 2.1291921490586256e-05, "grad_to_weight": 0.0005624036338593109, "update_to_weight": 0.00014847049907343345, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.04589042390929, "grad_norm": 0.03442581827078815, "grad_rms": 2.9563534048893515e-05, "grad_to_weight": 0.0008596592036378922, "update_to_weight": 0.0001767751226623339, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19274501679521, "grad_norm": 0.07661823778861102, "grad_rms": 0.00011943697523854117, "grad_to_weight": 0.0027176579557246865, "update_to_weight": 0.00013905056841299643, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403702020038642, "grad_norm": 0.007562228909753128, "grad_rms": 0.0006437396746519961, "grad_to_weight": 0.0008998687592350347, "update_to_weight": 8.05371039097305e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3034616982173997, "grad_norm": 0.07621601855234943, "grad_rms": 0.001764373142278593, "grad_to_weight": 0.058472004706069725, "update_to_weight": 0.000242923831420371, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41001}
{"step": 41050, "lr": 1.7658025948008258e-05, "window_steps": 49, "elapsed_seconds": 52.058344862, "steps_per_second": 0.9604608085897389, "eta_seconds": 9318.443730298, "images_seen": 42035200, "images_per_second": 983.5118679958927, "effective_batch_size": 1024, "loss": 0.3750106692314148, "edge_loss": 0.1956055909395218, "diag_loss": 0.17940503358840942, "jvp_rms": 2.1061861515045166, "grad_norm": 0.2826800048351288, "grad_norm_postclip": 0.2741435766220093, "clip_fraction": 0.020408162847161293, "ema_updates": 41050, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.38574159145355225, "updates_nfe_5": 10, "loss_nfe_4": 0.33549828827381134, "updates_nfe_4": 10, "loss_nfe_3": 0.37181579768657685, "updates_nfe_3": 10, "loss_nfe_1": 0.41067173264243384, "updates_nfe_1": 11, "loss_nfe_2": 0.36594706773757935, "updates_nfe_2": 8, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.476520825871905, "grad_norm": 0.06315738688705122, "grad_rms": 0.00018390814166854307, "grad_to_weight": 0.0038331749496459035, "update_to_weight": 6.263585874333062e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24529961316455, "grad_norm": 0.09837760579088599, "grad_rms": 0.00028646597895442264, "grad_to_weight": 0.006055758165959872, "update_to_weight": 9.335488573736456e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.057127185455498, "grad_norm": 0.004801248507161248, "grad_rms": 8.59561367484448e-05, "grad_to_weight": 0.000949402364443163, "update_to_weight": 3.4546737096461623e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.83367985509405, "grad_norm": 0.04232705162902778, "grad_rms": 0.00013544673858482522, "grad_to_weight": 0.003298122760341959, "update_to_weight": 9.505476613039358e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.060298299628747, "grad_norm": 0.0424993489501327, "grad_rms": 0.0009668936678883895, "grad_to_weight": 0.008398585702595971, "update_to_weight": 4.1933843925666844e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634273839338746, "grad_norm": 0.17737602956807738, "grad_rms": 0.00042741336030043534, "grad_to_weight": 0.01066328661420717, "update_to_weight": 9.503599108480481e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.5540747716265, "grad_norm": 0.02200891562686678, "grad_rms": 3.0364606954758916e-05, "grad_to_weight": 0.0008963447342882257, "update_to_weight": 0.000105267068822511, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.819135156536678, "grad_norm": 0.00967428210178324, "grad_rms": 1.1744353095974173e-05, "grad_to_weight": 0.0003139050480373828, "update_to_weight": 8.902022206839817e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.58679956430723, "grad_norm": 0.003098889815304489, "grad_rms": 2.7818803211834988e-06, "grad_to_weight": 8.959747228253149e-05, "update_to_weight": 0.00010118142099080184, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01204176145329, "grad_norm": 0.002458284936855627, "grad_rms": 1.9773505338661093e-06, "grad_to_weight": 6.641851732199877e-05, "update_to_weight": 0.00010306611872527368, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65561504054084, "grad_norm": 0.0009779803988063309, "grad_rms": 1.766908744518267e-06, "grad_to_weight": 3.966562574887112e-05, "update_to_weight": 6.459722124415042e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.25748404263901, "grad_norm": 0.002596179133078694, "grad_rms": 2.088267359833193e-06, "grad_to_weight": 6.968208401049087e-05, "update_to_weight": 0.00011032416760564887, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.35399561010833, "grad_norm": 0.0077574905493650965, "grad_rms": 3.4440570216810764e-06, "grad_to_weight": 0.00012244674506571083, "update_to_weight": 0.00012526098449736926, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.73787089510901, "grad_norm": 0.029210425023791578, "grad_rms": 2.1795972273072102e-05, "grad_to_weight": 0.000575712470950518, "update_to_weight": 0.00010218151853035229, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.04662962029779, "grad_norm": 0.06082983386201514, "grad_rms": 5.2238260552667715e-05, "grad_to_weight": 0.0015189751157281735, "update_to_weight": 0.0001318408279788061, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.193343611984684, "grad_norm": 0.2524884473405081, "grad_rms": 0.0003935937095842292, "grad_to_weight": 0.008955604940493047, "update_to_weight": 0.00011001360984028747, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.402647927141777, "grad_norm": 0.007002215583772388, "grad_rms": 0.0005960681745466929, "grad_to_weight": 0.0008333344017847293, "update_to_weight": 7.705260392546479e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3024817912463575, "grad_norm": 0.10330831714073402, "grad_rms": 0.0023915500127038784, "grad_to_weight": 0.07931651546688978, "update_to_weight": 0.000203713522876658, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41050}
{"step": 41100, "lr": 1.757518875626678e-05, "window_steps": 50, "elapsed_seconds": 79.56082202699999, "steps_per_second": 1.2569000351211017, "eta_seconds": 7080.913160402998, "images_seen": 42086400, "images_per_second": 1287.0656359640082, "effective_batch_size": 1024, "loss": 0.3781128227710724, "edge_loss": 0.19666515290737152, "diag_loss": 0.18144769966602325, "jvp_rms": 2.3402342796325684, "grad_norm": 0.28814345598220825, "grad_norm_postclip": 0.28054386377334595, "clip_fraction": 0.03999999910593033, "ema_updates": 41100, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39706082542737325, "updates_nfe_1": 15, "loss_nfe_5": 0.3687207115193208, "updates_nfe_5": 12, "loss_nfe_3": 0.3806685507297516, "updates_nfe_3": 9, "loss_nfe_4": 0.3529884107410908, "updates_nfe_4": 8, "loss_nfe_2": 0.379192719856898, "updates_nfe_2": 6, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.480112748732562, "grad_norm": 0.051869579862899584, "grad_rms": 0.00015103915015948386, "grad_to_weight": 0.003147404429431997, "update_to_weight": 5.407621854192982e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24332044208865, "grad_norm": 0.02213011326322038, "grad_rms": 6.444072824659036e-05, "grad_to_weight": 0.0013624131434283751, "update_to_weight": 5.021866094917782e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.056196644882273, "grad_norm": 0.0013698144375111786, "grad_rms": 2.4523612334392856e-05, "grad_to_weight": 0.00027091795151948107, "update_to_weight": 3.5322175827256704e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.833043148838897, "grad_norm": 0.016606843683365558, "grad_rms": 5.314196780853211e-05, "grad_to_weight": 0.0012940690287375919, "update_to_weight": 6.619770602978969e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.060595307142322, "grad_norm": 0.022807138005529064, "grad_rms": 0.0005188803561691698, "grad_to_weight": 0.004506809302324566, "update_to_weight": 3.363198527485811e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634333464394885, "grad_norm": 0.12666643490372753, "grad_rms": 0.0003052212111823141, "grad_to_weight": 0.007614758666155869, "update_to_weight": 9.12640336399522e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55392811226317, "grad_norm": 0.019350643628333063, "grad_rms": 2.669712120567529e-05, "grad_to_weight": 0.0007880874921462613, "update_to_weight": 0.00010877837930997101, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.81911733265673, "grad_norm": 0.006983033879508873, "grad_rms": 8.477240450429457e-06, "grad_to_weight": 0.0002265812419004444, "update_to_weight": 9.87503002545676e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.58668485893451, "grad_norm": 0.003551379821003401, "grad_rms": 3.1880816117777333e-06, "grad_to_weight": 0.00010268054991358908, "update_to_weight": 0.00011701840640603937, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01228911994552, "grad_norm": 0.002996467314646366, "grad_rms": 2.41024388812594e-06, "grad_to_weight": 8.0958713602818e-05, "update_to_weight": 0.00012388045715635464, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65631931232336, "grad_norm": 0.0010587675445583146, "grad_rms": 1.9128661833872636e-06, "grad_to_weight": 4.294102177810201e-05, "update_to_weight": 8.271789209428263e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.25809508568658, "grad_norm": 0.0031451749538453605, "grad_rms": 2.5298586347127268e-06, "grad_to_weight": 8.44158818804894e-05, "update_to_weight": 0.00012817073589577134, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.355920451446366, "grad_norm": 0.008750689785480995, "grad_rms": 3.885003070059208e-06, "grad_to_weight": 0.00013811952731690168, "update_to_weight": 0.00013553963206047387, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.73889820458881, "grad_norm": 0.023057161274403336, "grad_rms": 1.7204585261026593e-05, "grad_to_weight": 0.0004544277091203777, "update_to_weight": 9.994296292562568e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.047242304694805, "grad_norm": 0.0379347367466548, "grad_rms": 3.2576854749656634e-05, "grad_to_weight": 0.0009472496622372333, "update_to_weight": 0.00010576395251954589, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.193267841144277, "grad_norm": 0.10789648461364731, "grad_rms": 0.00016819532963783992, "grad_to_weight": 0.0038270300988729987, "update_to_weight": 7.555077775360422e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403337051933502, "grad_norm": 0.00959003246603867, "grad_rms": 0.0008163577766903845, "grad_to_weight": 0.0011412171625119006, "update_to_weight": 5.488077747177784e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3031076242582902, "grad_norm": 0.14782039459953653, "grad_rms": 0.0034219884358470745, "grad_to_weight": 0.11343682735620071, "update_to_weight": 6.217385428900782e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41100}
{"step": 41150, "lr": 1.7492760836265068e-05, "window_steps": 50, "elapsed_seconds": 107.05315034499999, "steps_per_second": 1.401173151061835, "eta_seconds": 6316.135870354999, "images_seen": 42137600, "images_per_second": 1434.801306687319, "effective_batch_size": 1024, "loss": 0.3469994366168976, "edge_loss": 0.18066911399364471, "diag_loss": 0.16633030772209167, "jvp_rms": 1.5847655534744263, "grad_norm": 0.24273787438869476, "grad_norm_postclip": 0.24273787438869476, "clip_fraction": 0.0, "ema_updates": 41150, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4128486752510071, "updates_nfe_1": 10, "loss_nfe_3": 0.2998237609863281, "updates_nfe_3": 9, "loss_nfe_5": 0.3178978455918176, "updates_nfe_5": 14, "loss_nfe_4": 0.36155744791030886, "updates_nfe_4": 5, "loss_nfe_2": 0.35539301733175915, "updates_nfe_2": 12, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47592718962464, "grad_norm": 0.21388645172450815, "grad_rms": 0.0006228164558973791, "grad_to_weight": 0.012981755094135067, "update_to_weight": 0.00011746332907717389, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24713766858258, "grad_norm": 0.05041547779695929, "grad_rms": 0.00014680494697400824, "grad_to_weight": 0.0031030375211535703, "update_to_weight": 0.0001128549641007165, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.055515321750794, "grad_norm": 0.006640836034342147, "grad_rms": 0.00011889003650623634, "grad_to_weight": 0.0013135824167658412, "update_to_weight": 5.7410563064721585e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.834373002647762, "grad_norm": 0.039581491474279855, "grad_rms": 0.0001266609348437959, "grad_to_weight": 0.0030840222164428364, "update_to_weight": 0.0001426937910305244, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.059796965675827, "grad_norm": 0.09527596688129046, "grad_rms": 0.0021676033011130635, "grad_to_weight": 0.018829998027117406, "update_to_weight": 5.651409869750872e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634238981514624, "grad_norm": 0.22091258920253218, "grad_rms": 0.0005323210374797849, "grad_to_weight": 0.01328059488913373, "update_to_weight": 0.00016776894269252555, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.554929852844722, "grad_norm": 0.06393757085828873, "grad_rms": 8.821148854712435e-05, "grad_to_weight": 0.002603858827594308, "update_to_weight": 0.00020515627239605126, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.81971046623664, "grad_norm": 0.04203933464349276, "grad_rms": 5.103477289358672e-05, "grad_to_weight": 0.0013640405444284543, "update_to_weight": 0.0001665749734823893, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.58587485170518, "grad_norm": 0.028103176895657004, "grad_rms": 2.5228284782072913e-05, "grad_to_weight": 0.0008125622675776102, "update_to_weight": 0.0001875884052242573, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.011589915672864, "grad_norm": 0.02237264887790017, "grad_rms": 1.7995704460240528e-05, "grad_to_weight": 0.0006044768389813561, "update_to_weight": 0.00018155982993132454, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65681439468104, "grad_norm": 0.007823816221206844, "grad_rms": 1.4135221231046157e-05, "grad_to_weight": 0.0003173084769172207, "update_to_weight": 0.00012945866439750033, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.25866680321674, "grad_norm": 0.018432516925109613, "grad_rms": 1.4826412770922055e-05, "grad_to_weight": 0.0004947175652435915, "update_to_weight": 0.00018262586199646225, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.356666094988576, "grad_norm": 0.035767722844589456, "grad_rms": 1.5879629659688466e-05, "grad_to_weight": 0.000564545533235036, "update_to_weight": 0.00019799144539191287, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.73911713670726, "grad_norm": 0.04637757261356104, "grad_rms": 3.460560008812754e-05, "grad_to_weight": 0.000914039802636005, "update_to_weight": 0.00015257128372836887, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.04797233182505, "grad_norm": 0.04997552527874506, "grad_rms": 4.291700872781296e-05, "grad_to_weight": 0.0012478915253102802, "update_to_weight": 0.00016920619251166856, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19359365431323, "grad_norm": 0.12619065419563255, "grad_rms": 0.00019671334757244237, "grad_to_weight": 0.0044758626992670416, "update_to_weight": 0.00012370120003409653, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.402892168876011, "grad_norm": 0.006887468397623266, "grad_rms": 0.0005863002453871299, "grad_to_weight": 0.0008196545021884469, "update_to_weight": 3.8652653416472615e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3027819577610726, "grad_norm": 0.08820588473318236, "grad_rms": 0.002041934188772328, "grad_to_weight": 0.0677057923681801, "update_to_weight": 0.00011349425136901395, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41150}
{"step": 41200, "lr": 1.7410743089421665e-05, "window_steps": 50, "elapsed_seconds": 134.52829757299997, "steps_per_second": 1.486676064502137, "eta_seconds": 5919.245093211999, "images_seen": 42188800, "images_per_second": 1522.3562900501884, "effective_batch_size": 1024, "loss": 0.3830133080482483, "edge_loss": 0.19913604855537415, "diag_loss": 0.18387721478939056, "jvp_rms": 2.303164005279541, "grad_norm": 0.28074827790260315, "grad_norm_postclip": 0.28074827790260315, "clip_fraction": 0.0, "ema_updates": 41200, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.41501879394054414, "updates_nfe_1": 10, "loss_nfe_5": 0.36887693256139753, "updates_nfe_5": 10, "loss_nfe_2": 0.35235551541501825, "updates_nfe_2": 11, "loss_nfe_3": 0.3498629050122367, "updates_nfe_3": 9, "loss_nfe_4": 0.42870292663574217, "updates_nfe_4": 10, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47992664357868, "grad_norm": 0.07894480148274614, "grad_rms": 0.0002298795509232911, "grad_to_weight": 0.0047903612188411404, "update_to_weight": 7.979665446487964e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24045411475663, "grad_norm": 0.05853487474197549, "grad_rms": 0.00017044783780952697, "grad_to_weight": 0.003604263423200014, "update_to_weight": 7.44813658201559e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.056230972688048, "grad_norm": 0.005848464020817156, "grad_rms": 0.00010470430188979083, "grad_to_weight": 0.001156684505199321, "update_to_weight": 4.401127767448352e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.832320202226132, "grad_norm": 0.041878932881175585, "grad_rms": 0.0001340127567562003, "grad_to_weight": 0.003263551113220389, "update_to_weight": 0.00011250896184736672, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0608069332059635, "grad_norm": 0.10029597756126479, "grad_rms": 0.0022818124986444086, "grad_to_weight": 0.01981817897521106, "update_to_weight": 4.843914445950937e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63463617292151, "grad_norm": 0.18751705608675298, "grad_rms": 0.00045184963972216754, "grad_to_weight": 0.011272687550089033, "update_to_weight": 0.00012256439146641267, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.554058614281953, "grad_norm": 0.06686142290076491, "grad_rms": 9.224538188239074e-05, "grad_to_weight": 0.002723029375757649, "update_to_weight": 0.00015468489725261564, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8200550527904, "grad_norm": 0.04832740826033638, "grad_rms": 5.866834777518733e-05, "grad_to_weight": 0.0015680506792592796, "update_to_weight": 0.00014011635550212313, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.5874472013511, "grad_norm": 0.034616089001709135, "grad_rms": 3.1074940552776384e-05, "grad_to_weight": 0.001000828098130264, "update_to_weight": 0.0001672897739346536, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0127937261437, "grad_norm": 0.02990392257968684, "grad_rms": 2.4053573445097852e-05, "grad_to_weight": 0.0008079347590172432, "update_to_weight": 0.00017429094601840542, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65740600506139, "grad_norm": 0.010534950752043332, "grad_rms": 1.9033404585178028e-05, "grad_to_weight": 0.0004272530026021731, "update_to_weight": 0.00011806322968916661, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.260219733777404, "grad_norm": 0.02643902361892146, "grad_rms": 2.1266541027844784e-05, "grad_to_weight": 0.0007095777697455113, "update_to_weight": 0.00016793781410141405, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.359080224345114, "grad_norm": 0.04089566979703093, "grad_rms": 1.8156260433000795e-05, "grad_to_weight": 0.0006454587038231209, "update_to_weight": 0.00016871941945329513, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7394587651588, "grad_norm": 0.051914723982661844, "grad_rms": 3.873726190456552e-05, "grad_to_weight": 0.0010231627464325666, "update_to_weight": 0.00012007850845537432, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.04874349425647, "grad_norm": 0.059479665969513845, "grad_rms": 5.107878965359674e-05, "grad_to_weight": 0.0014851818254433883, "update_to_weight": 0.00015021852861423307, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.193574170575086, "grad_norm": 0.14506277730365497, "grad_rms": 0.00022613231315307108, "grad_to_weight": 0.005145242544489208, "update_to_weight": 0.0001212068984581322, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403869972672002, "grad_norm": 0.007569871451471591, "grad_rms": 0.00064439025100694, "grad_to_weight": 0.0009007601826405648, "update_to_weight": 5.100502071659685e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3038201561372893, "grad_norm": 0.0627172427281418, "grad_rms": 0.001451881385686601, "grad_to_weight": 0.04810267921762196, "update_to_weight": 0.00015763214796127862, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41200}
{"step": 41250, "lr": 1.73291364126695e-05, "window_steps": 50, "elapsed_seconds": 161.996338653, "steps_per_second": 1.5432447552750308, "eta_seconds": 5669.871852855, "images_seen": 42240000, "images_per_second": 1580.2826294016315, "effective_batch_size": 1024, "loss": 0.34203147888183594, "edge_loss": 0.17816387116909027, "diag_loss": 0.16386760771274567, "jvp_rms": 1.4647655487060547, "grad_norm": 0.23330773413181305, "grad_norm_postclip": 0.23330773413181305, "clip_fraction": 0.0, "ema_updates": 41250, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3567957766354084, "updates_nfe_2": 16, "loss_nfe_5": 0.29225707550843555, "updates_nfe_5": 6, "loss_nfe_3": 0.33665566552769055, "updates_nfe_3": 11, "loss_nfe_4": 0.2856990061700344, "updates_nfe_4": 8, "loss_nfe_1": 0.40561073356204563, "updates_nfe_1": 9, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47355428323397, "grad_norm": 0.03144598391509144, "grad_rms": 9.156763364997693e-05, "grad_to_weight": 0.0019088766986427284, "update_to_weight": 5.4807086730024196e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248954856714533, "grad_norm": 0.019951600449674086, "grad_rms": 5.8097111721465415e-05, "grad_to_weight": 0.0012278697692011565, "update_to_weight": 5.6411259082309194e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.056430899233243, "grad_norm": 0.0014092944596449798, "grad_rms": 2.5230418111328384e-05, "grad_to_weight": 0.0002787132836836918, "update_to_weight": 3.314850932956871e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.834468114314038, "grad_norm": 0.013069592639406205, "grad_rms": 4.182274997925409e-05, "grad_to_weight": 0.0010183197716491217, "update_to_weight": 9.318508692460975e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.059414335289975, "grad_norm": 0.049420278048968404, "grad_rms": 0.0011243502569156867, "grad_to_weight": 0.009767983955031414, "update_to_weight": 4.8774470477809954e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.633702345266897, "grad_norm": 0.08406011913962468, "grad_rms": 0.00020255509200544204, "grad_to_weight": 0.005053602462926355, "update_to_weight": 0.00011587462627502405, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55448118748834, "grad_norm": 0.02526250596373451, "grad_rms": 3.4853423825417924e-05, "grad_to_weight": 0.001028834849771004, "update_to_weight": 0.0001261629927519124, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.821064035513928, "grad_norm": 0.010673746879688119, "grad_rms": 1.2957680052456107e-05, "grad_to_weight": 0.00034631338059546454, "update_to_weight": 9.40678760766938e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.5879289507418, "grad_norm": 0.00688877531099509, "grad_rms": 6.18406901080121e-06, "grad_to_weight": 0.00019916703659261299, "update_to_weight": 0.00010188232469519903, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.014368514976724, "grad_norm": 0.005678926443408499, "grad_rms": 4.567911581894725e-06, "grad_to_weight": 0.0001534249177075842, "update_to_weight": 0.00010201792153676714, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65759289158273, "grad_norm": 0.0021084320006801095, "grad_rms": 3.8092858954748554e-06, "grad_to_weight": 8.550842776708577e-05, "update_to_weight": 6.364932535859067e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26110756162181, "grad_norm": 0.005736715585819633, "grad_rms": 4.614394961378179e-06, "grad_to_weight": 0.0001539598783082963, "update_to_weight": 0.00010883712729243064, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.361561691612366, "grad_norm": 0.014772914461842975, "grad_rms": 6.5586621677764515e-06, "grad_to_weight": 0.00023315262546312165, "update_to_weight": 0.00012380167633437684, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7398990294189, "grad_norm": 0.03253073986017175, "grad_rms": 2.4273494940156472e-05, "grad_to_weight": 0.0006411274062904714, "update_to_weight": 0.00010672267272847506, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.04948568597277, "grad_norm": 0.039778498320506325, "grad_rms": 3.416020442163389e-05, "grad_to_weight": 0.0009932336867546502, "update_to_weight": 0.00013577816311302209, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.193568758423215, "grad_norm": 0.10072716732993024, "grad_rms": 0.0001570193799474407, "grad_to_weight": 0.003572700149917581, "update_to_weight": 9.956358680009886e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403066039242063, "grad_norm": 0.013312714135099914, "grad_rms": 0.0011332534849628107, "grad_to_weight": 0.0015842686553848254, "update_to_weight": 7.400141694718651e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3031349767216123, "grad_norm": 0.11963149865903203, "grad_rms": 0.00276942573508423, "grad_to_weight": 0.09180284528928642, "update_to_weight": 0.0002281194942328937, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41250}
{"step": 41300, "lr": 1.724794169844612e-05, "window_steps": 50, "elapsed_seconds": 189.47849858999996, "steps_per_second": 1.5832931030826365, "eta_seconds": 5494.876459109999, "images_seen": 42291200, "images_per_second": 1621.2921375566198, "effective_batch_size": 1024, "loss": 0.3994542956352234, "edge_loss": 0.20783250033855438, "diag_loss": 0.1916218400001526, "jvp_rms": 2.3730859756469727, "grad_norm": 0.25971469283103943, "grad_norm_postclip": 0.25971469283103943, "clip_fraction": 0.0, "ema_updates": 41300, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3873378742824901, "updates_nfe_4": 11, "loss_nfe_5": 0.4733388694849881, "updates_nfe_5": 11, "loss_nfe_1": 0.41733779311180114, "updates_nfe_1": 10, "loss_nfe_3": 0.35733024369586597, "updates_nfe_3": 11, "loss_nfe_2": 0.3430370305265699, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.478588661573653, "grad_norm": 0.11953454091023495, "grad_rms": 0.00034807290750198637, "grad_to_weight": 0.0072539307440191774, "update_to_weight": 8.56074767899339e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.241702734468568, "grad_norm": 0.11737133547614893, "grad_rms": 0.00034177386457068964, "grad_to_weight": 0.007226541292807952, "update_to_weight": 8.527446821535918e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.05620475531895, "grad_norm": 0.004613841057871539, "grad_rms": 8.260100519989102e-05, "grad_to_weight": 0.0009125107231897483, "update_to_weight": 5.574456594535778e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.832567530650714, "grad_norm": 0.04691863512565288, "grad_rms": 0.00015013982458118765, "grad_to_weight": 0.003656215719386417, "update_to_weight": 0.00013059079597529025, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.060149223931744, "grad_norm": 0.032648722528823226, "grad_rms": 0.0007427841568774341, "grad_to_weight": 0.006452126426314186, "update_to_weight": 4.561489222801974e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634789359682784, "grad_norm": 0.1339857924361365, "grad_rms": 0.0003228582684881128, "grad_to_weight": 0.008054552993672024, "update_to_weight": 0.0001215901234600162, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.554952223755954, "grad_norm": 0.019135420381157252, "grad_rms": 2.640018838904698e-05, "grad_to_weight": 0.0007792896604638669, "update_to_weight": 0.00017133559960338859, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.820784811019777, "grad_norm": 0.008572856761909373, "grad_rms": 1.0407248392572184e-05, "grad_to_weight": 0.0002781517996532068, "update_to_weight": 0.00015331878037774292, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.58767307735879, "grad_norm": 0.0028185773028827177, "grad_rms": 2.5302431515634453e-06, "grad_to_weight": 8.149080444292069e-05, "update_to_weight": 0.00017853502948374163, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0136347211867, "grad_norm": 0.0022426593950518265, "grad_rms": 1.8039095816767843e-06, "grad_to_weight": 6.05900882727446e-05, "update_to_weight": 0.00018081214120819957, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65788497628915, "grad_norm": 0.0009574210061565027, "grad_rms": 1.7297642672881464e-06, "grad_to_weight": 3.882818851159181e-05, "update_to_weight": 0.00012089902983433086, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26140076949184, "grad_norm": 0.002320116156551663, "grad_rms": 1.8662128429493143e-06, "grad_to_weight": 6.226594031996999e-05, "update_to_weight": 0.0001746354002643977, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.36142105204377, "grad_norm": 0.0068625396585415095, "grad_rms": 3.046729834495141e-06, "grad_to_weight": 0.00010830785586239868, "update_to_weight": 0.00018160598793383356, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74046439048917, "grad_norm": 0.02260446173771784, "grad_rms": 1.6866793991587922e-05, "grad_to_weight": 0.00044549181820170404, "update_to_weight": 0.00013552434043472047, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05043817058997, "grad_norm": 0.0557905942439227, "grad_rms": 4.7910760452069574e-05, "grad_to_weight": 0.0013930083362955843, "update_to_weight": 0.0001589881644422405, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.194456337444468, "grad_norm": 0.1803504490048557, "grad_rms": 0.00028114079276376466, "grad_to_weight": 0.00639666347335579, "update_to_weight": 0.0001222474443377409, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403687948179126, "grad_norm": 0.012020237594224449, "grad_rms": 0.0010232305753355392, "grad_to_weight": 0.001430352681863793, "update_to_weight": 3.5998791656112323e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3037095202034665, "grad_norm": 0.17423131345772444, "grad_rms": 0.004033391613045045, "grad_to_weight": 0.13364274077751048, "update_to_weight": 0.00010850999419922735, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41300}
{"step": 41350, "lr": 1.7167159834683903e-05, "window_steps": 50, "elapsed_seconds": 216.93979584299996, "steps_per_second": 1.613350831459693, "eta_seconds": 5361.512097262713, "images_seen": 42342400, "images_per_second": 1652.0712514147256, "effective_batch_size": 1024, "loss": 0.3739354610443115, "edge_loss": 0.1925373375415802, "diag_loss": 0.18139801919460297, "jvp_rms": 2.028437376022339, "grad_norm": 0.3756810426712036, "grad_norm_postclip": 0.3461199104785919, "clip_fraction": 0.07999999821186066, "ema_updates": 41350, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.37652163666028243, "updates_nfe_3": 13, "loss_nfe_2": 0.351807747568403, "updates_nfe_2": 7, "loss_nfe_4": 0.35135468201977865, "updates_nfe_4": 14, "loss_nfe_5": 0.3950935259461403, "updates_nfe_5": 10, "loss_nfe_1": 0.41157230238119763, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47406001325455, "grad_norm": 0.09353975722410902, "grad_rms": 0.0002723786364685639, "grad_to_weight": 0.0056780026993254636, "update_to_weight": 6.276085745418104e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24895016139985, "grad_norm": 0.07763454746113092, "grad_rms": 0.00022606421919242124, "grad_to_weight": 0.004777819286168743, "update_to_weight": 8.930742347458831e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.056608940370026, "grad_norm": 0.0034647367230162584, "grad_rms": 6.202873754956463e-05, "grad_to_weight": 0.0006851897712229889, "update_to_weight": 3.394179660382176e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.834537664273851, "grad_norm": 0.03143884550849831, "grad_rms": 0.00010060443440095305, "grad_to_weight": 0.0024495502939705658, "update_to_weight": 8.62133438244405e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.059312546964851, "grad_norm": 0.04505456370258472, "grad_rms": 0.0010250268164017891, "grad_to_weight": 0.008905273845873298, "update_to_weight": 2.7707954056522404e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634191281215873, "grad_norm": 0.12463985656256557, "grad_rms": 0.00030033787570109185, "grad_to_weight": 0.007492991661296746, "update_to_weight": 0.00010447962551004472, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.554832911993874, "grad_norm": 0.03651649159156874, "grad_rms": 5.037998842574486e-05, "grad_to_weight": 0.00148714070759293, "update_to_weight": 0.00013085193984374861, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82145019288008, "grad_norm": 0.02163209107058967, "grad_rms": 2.6260854610642957e-05, "grad_to_weight": 0.0007018518251158344, "update_to_weight": 0.00010787761625916517, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.58813894196607, "grad_norm": 0.013814584727046742, "grad_rms": 1.240138362057812e-05, "grad_to_weight": 0.00039940237172707185, "update_to_weight": 0.0001165790477146357, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.015056124269385, "grad_norm": 0.01128515696636982, "grad_rms": 9.077349341267483e-06, "grad_to_weight": 0.0003048801797971763, "update_to_weight": 0.00012034642781773705, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65786269876393, "grad_norm": 0.004498028747542357, "grad_rms": 8.126549710840371e-06, "grad_to_weight": 0.000182417624856344, "update_to_weight": 8.099969833056596e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.2629486722401, "grad_norm": 0.009806053104039985, "grad_rms": 7.88761467382804e-06, "grad_to_weight": 0.00026315826990216777, "update_to_weight": 0.00011955841105135831, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.36341694875581, "grad_norm": 0.020025317879320246, "grad_rms": 8.89054729938577e-06, "grad_to_weight": 0.00031603910968872485, "update_to_weight": 0.00013100161278369542, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7409238920802, "grad_norm": 0.03294811422738745, "grad_rms": 2.4584927592297444e-05, "grad_to_weight": 0.0006493400533554355, "update_to_weight": 0.00010071543916739792, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05112242155721, "grad_norm": 0.03850278547750859, "grad_rms": 3.306467258056283e-05, "grad_to_weight": 0.0009613409849603806, "update_to_weight": 0.00011887044703327721, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19402554038488, "grad_norm": 0.09507829705303489, "grad_rms": 0.00014821359168004766, "grad_to_weight": 0.003372285270751619, "update_to_weight": 8.95212190284966e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403434196733388, "grad_norm": 0.004525192369267528, "grad_rms": 0.0003852099557278638, "grad_to_weight": 0.0005384932235236133, "update_to_weight": 7.5114396342318916e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.303576424620847, "grad_norm": 0.06772537506044722, "grad_rms": 0.001567817829861096, "grad_to_weight": 0.05195351325883755, "update_to_weight": 0.0002772201919293189, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41350}
{"step": 41400, "lr": 1.7086791704800387e-05, "window_steps": 50, "elapsed_seconds": 244.44485288599998, "steps_per_second": 1.636360902172668, "eta_seconds": 5255.564337049, "images_seen": 42393600, "images_per_second": 1675.633563824812, "effective_batch_size": 1024, "loss": 0.36449697613716125, "edge_loss": 0.18949614465236664, "diag_loss": 0.17500080168247223, "jvp_rms": 1.8651561737060547, "grad_norm": 0.23135097324848175, "grad_norm_postclip": 0.22902525961399078, "clip_fraction": 0.019999999552965164, "ema_updates": 41400, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4000436273904947, "updates_nfe_1": 13, "loss_nfe_5": 0.3702174375454585, "updates_nfe_5": 6, "loss_nfe_4": 0.3430156434575717, "updates_nfe_4": 12, "loss_nfe_2": 0.36029690876603127, "updates_nfe_2": 8, "loss_nfe_3": 0.34585572237318213, "updates_nfe_3": 11, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47898497483329, "grad_norm": 0.05149261441963193, "grad_rms": 0.0001499414636090793, "grad_to_weight": 0.0031247443030181453, "update_to_weight": 4.196427393422074e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.242491880264826, "grad_norm": 0.019666884758463277, "grad_rms": 5.7268047438484606e-05, "grad_to_weight": 0.0012108292805957444, "update_to_weight": 4.8283066636015795e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0572818188975734, "grad_norm": 0.004076415466478307, "grad_rms": 7.297954370773773e-05, "grad_to_weight": 0.0008060487060946342, "update_to_weight": 3.574923879359704e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.833868899048671, "grad_norm": 0.020150948292043544, "grad_rms": 6.448311707298202e-05, "grad_to_weight": 0.0015701382373897604, "update_to_weight": 7.219404864098667e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.059768881966449, "grad_norm": 0.03577349469417108, "grad_rms": 0.0008138751852085704, "grad_to_weight": 0.007070183545669762, "update_to_weight": 2.481058857194225e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63480036706041, "grad_norm": 0.0498983960167178, "grad_rms": 0.00012023744790680277, "grad_to_weight": 0.0029996390047171634, "update_to_weight": 7.61000923099627e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.553764051598332, "grad_norm": 0.025718567881020473, "grad_rms": 3.548262978450285e-05, "grad_to_weight": 0.0010474389110758894, "update_to_weight": 9.9930379735713e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.820902640104187, "grad_norm": 0.01306541944046896, "grad_rms": 1.5861114823971135e-05, "grad_to_weight": 0.00042391423745871166, "update_to_weight": 8.408343460024452e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.590519341581334, "grad_norm": 0.008376426358969818, "grad_rms": 7.519536685292331e-06, "grad_to_weight": 0.00024215960090834772, "update_to_weight": 9.911383275272894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0160157878274, "grad_norm": 0.007525199595629243, "grad_rms": 6.052983205803371e-06, "grad_to_weight": 0.0002032957744226995, "update_to_weight": 0.00010528366115259478, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.658696854586747, "grad_norm": 0.0030955643979567655, "grad_rms": 5.592729476628675e-06, "grad_to_weight": 0.00012553641484833622, "update_to_weight": 6.989986152107043e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26465047205847, "grad_norm": 0.007078338329127868, "grad_rms": 5.693545066378091e-06, "grad_to_weight": 0.0001899477987707224, "update_to_weight": 0.00011237974971477239, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.36514885837285, "grad_norm": 0.018582772971183947, "grad_rms": 8.25010734160023e-06, "grad_to_weight": 0.00029326488307820784, "update_to_weight": 0.00012775526863979702, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74191264840881, "grad_norm": 0.02942865061243462, "grad_rms": 2.1958805880438806e-05, "grad_to_weight": 0.0005799673105810185, "update_to_weight": 9.696508890199706e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0515079743026, "grad_norm": 0.040970669814534205, "grad_rms": 3.518399424933191e-05, "grad_to_weight": 0.001022949493957165, "update_to_weight": 0.00010927698627887434, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.194291812711178, "grad_norm": 0.09782895076301448, "grad_rms": 0.00015250147102223566, "grad_to_weight": 0.0034698140819734674, "update_to_weight": 7.823155714531769e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403993438413112, "grad_norm": 0.013886299738953816, "grad_rms": 0.0011820803340857943, "grad_to_weight": 0.0016523453808854833, "update_to_weight": 5.384232968797986e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3041158097479995, "grad_norm": 0.16127299944102932, "grad_rms": 0.0037334113509618814, "grad_to_weight": 0.1236646302694489, "update_to_weight": 0.00015896378286188133, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41400}
{"step": 41450, "lr": 1.700683818768856e-05, "window_steps": 50, "elapsed_seconds": 271.902368969, "steps_per_second": 1.6550058085419077, "eta_seconds": 5166.145010411001, "images_seen": 42444800, "images_per_second": 1694.7259479469135, "effective_batch_size": 1024, "loss": 0.35306864976882935, "edge_loss": 0.18201249837875366, "diag_loss": 0.17105616629123688, "jvp_rms": 1.8180468082427979, "grad_norm": 0.18766821920871735, "grad_norm_postclip": 0.1862618625164032, "clip_fraction": 0.019999999552965164, "ema_updates": 41450, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3040771550602383, "updates_nfe_4": 9, "loss_nfe_5": 0.33976279542996335, "updates_nfe_5": 13, "loss_nfe_2": 0.3717987537384033, "updates_nfe_2": 10, "loss_nfe_1": 0.4049018360674381, "updates_nfe_1": 8, "loss_nfe_3": 0.354261988401413, "updates_nfe_3": 10, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.478808116588997, "grad_norm": 0.014549819463267506, "grad_rms": 4.236765311217991e-05, "grad_to_weight": 0.0008829412515957629, "update_to_weight": 0.0001337712960643, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245226349620083, "grad_norm": 0.009237885086732043, "grad_rms": 2.6899819055002062e-05, "grad_to_weight": 0.0005686522851648714, "update_to_weight": 0.00010958199694101052, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.056453909139676, "grad_norm": 0.0024771077678104997, "grad_rms": 4.434734292819292e-05, "grad_to_weight": 0.000489890308963968, "update_to_weight": 6.424346975989937e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.83440748070823, "grad_norm": 0.010895546383651408, "grad_rms": 3.486579305592818e-05, "grad_to_weight": 0.0008489325588289775, "update_to_weight": 0.0001342598997041718, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.059413769804936, "grad_norm": 0.043021393356698016, "grad_rms": 0.0009787705893832771, "grad_to_weight": 0.00850323680056646, "update_to_weight": 2.964810159834506e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634814126272197, "grad_norm": 0.054206471774078516, "grad_rms": 0.0001306183834839839, "grad_to_weight": 0.003258616018345977, "update_to_weight": 0.0001019252397629145, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.554523444409, "grad_norm": 0.022780839335095043, "grad_rms": 3.142959172714809e-05, "grad_to_weight": 0.0009277654842974434, "update_to_weight": 0.00011078637797558394, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.821522473080297, "grad_norm": 0.012168444801933034, "grad_rms": 1.4772208509035587e-05, "grad_to_weight": 0.0003948034952705866, "update_to_weight": 0.00010082692480155584, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.589497677898926, "grad_norm": 0.009045383659895418, "grad_rms": 8.120061151172413e-06, "grad_to_weight": 0.000261506649912265, "update_to_weight": 0.00012025354626844564, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01598445899757, "grad_norm": 0.007396108260506416, "grad_rms": 5.949147065169035e-06, "grad_to_weight": 0.00019980849810165253, "update_to_weight": 0.0001272512213316215, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.658164681281914, "grad_norm": 0.002973387066782512, "grad_rms": 5.371992747040659e-06, "grad_to_weight": 0.00012058428132080806, "update_to_weight": 8.470602578429049e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26426229164647, "grad_norm": 0.007402312413853243, "grad_rms": 5.954137449216308e-06, "grad_to_weight": 0.0001986437395679404, "update_to_weight": 0.00012556506443953203, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.36549369345768, "grad_norm": 0.01827418531642535, "grad_rms": 8.113105114860499e-06, "grad_to_weight": 0.00028839332342031626, "update_to_weight": 0.00013108451544040597, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74236973119702, "grad_norm": 0.02518108334838566, "grad_rms": 1.8789394335080786e-05, "grad_to_weight": 0.0004962535940236947, "update_to_weight": 9.933423887462331e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.051855425771166, "grad_norm": 0.030965171986068893, "grad_rms": 2.65916676007316e-05, "grad_to_weight": 0.0007731270288702907, "update_to_weight": 0.00012042216354937869, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.194353509598667, "grad_norm": 0.08501850841335479, "grad_rms": 0.00013253180675075463, "grad_to_weight": 0.003015444506802064, "update_to_weight": 0.0001022365586490189, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.403824580399391, "grad_norm": 0.012633085770554518, "grad_rms": 0.0010753996765819952, "grad_to_weight": 0.0015032543396990008, "update_to_weight": 6.114498883558759e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3040520953876982, "grad_norm": 0.15740598783003554, "grad_rms": 0.0036438914369476038, "grad_to_weight": 0.12070529113581027, "update_to_weight": 0.0001153680195950622, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41450}
{"step": 41500, "lr": 1.692730015770726e-05, "window_steps": 50, "elapsed_seconds": 299.358317682, "steps_per_second": 1.6702392098927281, "eta_seconds": 5089.091400594, "images_seen": 42496000, "images_per_second": 1710.3249509301536, "effective_batch_size": 1024, "loss": 0.3659754991531372, "edge_loss": 0.18785573542118073, "diag_loss": 0.17811976373195648, "jvp_rms": 2.0639452934265137, "grad_norm": 0.24767707288265228, "grad_norm_postclip": 0.24688956141471863, "clip_fraction": 0.019999999552965164, "ema_updates": 41500, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3842547684907913, "updates_nfe_5": 6, "loss_nfe_3": 0.37228949446427195, "updates_nfe_3": 19, "loss_nfe_4": 0.32660318504680286, "updates_nfe_4": 11, "loss_nfe_2": 0.35702450360570637, "updates_nfe_2": 7, "loss_nfe_1": 0.4039912053516933, "updates_nfe_1": 7, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.48177926837884, "grad_norm": 0.05348083886518682, "grad_rms": 0.00015573097899317662, "grad_to_weight": 0.0032448462022418066, "update_to_weight": 4.600706938044473e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245062914368688, "grad_norm": 0.014718019111818624, "grad_rms": 4.285743405973016e-05, "grad_to_weight": 0.0009059995144001936, "update_to_weight": 4.78254625413549e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.057018562064069, "grad_norm": 0.001542523396709343, "grad_rms": 2.7615598698434733e-05, "grad_to_weight": 0.0003050262477343464, "update_to_weight": 3.2326480758832336e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.835741948259349, "grad_norm": 0.017130635318621735, "grad_rms": 5.4818103186806534e-05, "grad_to_weight": 0.0013346042159210606, "update_to_weight": 6.607323722066541e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.059450149213641, "grad_norm": 0.01576409048484801, "grad_rms": 0.0003586454768449238, "grad_to_weight": 0.0031157714810764814, "update_to_weight": 3.1356247962513365e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634930620475906, "grad_norm": 0.06239670360847995, "grad_rms": 0.00015035393917606557, "grad_to_weight": 0.0037509446256226617, "update_to_weight": 8.899914502324896e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.554739699276364, "grad_norm": 0.013560544839393488, "grad_rms": 1.8708809698825967e-05, "grad_to_weight": 0.0005522577313166599, "update_to_weight": 0.00010984424466600707, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.821847236819597, "grad_norm": 0.0068528849234797224, "grad_rms": 8.319242649807432e-06, "grad_to_weight": 0.0002223385532614446, "update_to_weight": 0.00010132369936274995, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.589827648570505, "grad_norm": 0.0036178737237449886, "grad_rms": 3.2477733370536564e-06, "grad_to_weight": 0.00010459357474984425, "update_to_weight": 0.00012139194113706066, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01697707534848, "grad_norm": 0.0030481809277435046, "grad_rms": 2.451840343822637e-06, "grad_to_weight": 8.234548492544104e-05, "update_to_weight": 0.0001297743094059713, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.657240157558185, "grad_norm": 0.001086179993116, "grad_rms": 1.9623920175699822e-06, "grad_to_weight": 4.405115844982567e-05, "update_to_weight": 8.410629018994663e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.264806070666395, "grad_norm": 0.003192806327929436, "grad_rms": 2.5681714932271646e-06, "grad_to_weight": 8.567886605594617e-05, "update_to_weight": 0.0001268224311142564, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.3675434007081, "grad_norm": 0.008758794136356378, "grad_rms": 3.888601121047593e-06, "grad_to_weight": 0.00013822208762251785, "update_to_weight": 0.00013476159888940848, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74290619679637, "grad_norm": 0.02000200781089765, "grad_rms": 1.4924918322722434e-05, "grad_to_weight": 0.00039418333142614655, "update_to_weight": 9.799328210266863e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05250003169668, "grad_norm": 0.029840906724562195, "grad_rms": 2.562619296547087e-05, "grad_to_weight": 0.0007450447962286187, "update_to_weight": 0.00011622733849202149, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19473451194777, "grad_norm": 0.06011769828988372, "grad_rms": 9.371497243067954e-05, "grad_to_weight": 0.0021322314017324156, "update_to_weight": 8.652682523232662e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.404284846682655, "grad_norm": 0.006579145237793412, "grad_rms": 0.0005600540350481963, "grad_to_weight": 0.0007828322525729642, "update_to_weight": 6.9423902728251135e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3045859382982705, "grad_norm": 0.06349330461232086, "grad_rms": 0.001469846936383154, "grad_to_weight": 0.04866931548805659, "update_to_weight": 0.0001170181175593809, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41500}
{"step": 41550, "lr": 1.684817848467165e-05, "window_steps": 50, "elapsed_seconds": 328.79207246299995, "steps_per_second": 1.6727897235475264, "eta_seconds": 5051.441840567909, "images_seen": 42547200, "images_per_second": 1712.936676912667, "effective_batch_size": 1024, "loss": 0.3818948566913605, "edge_loss": 0.19644774496555328, "diag_loss": 0.1854470819234848, "jvp_rms": 2.3954687118530273, "grad_norm": 0.25354039669036865, "grad_norm_postclip": 0.25354039669036865, "clip_fraction": 0.0, "ema_updates": 41550, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3853680044412613, "updates_nfe_5": 13, "loss_nfe_3": 0.3593549933284521, "updates_nfe_3": 8, "loss_nfe_2": 0.39239485189318657, "updates_nfe_2": 8, "loss_nfe_4": 0.37322743485371274, "updates_nfe_4": 12, "loss_nfe_1": 0.3991367816925049, "updates_nfe_1": 9, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.481976462054032, "grad_norm": 0.03131840899754665, "grad_rms": 9.119614795106512e-05, "grad_to_weight": 0.00190016100736766, "update_to_weight": 6.320364080896568e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.246261401365064, "grad_norm": 0.0669893164061176, "grad_rms": 0.00019506634614138864, "grad_to_weight": 0.004123368124588279, "update_to_weight": 6.762258777513677e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0570102643479675, "grad_norm": 0.002588713930799459, "grad_rms": 4.6345413761964494e-05, "grad_to_weight": 0.000511906006806027, "update_to_weight": 3.303903112480823e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.836032600073757, "grad_norm": 0.027842466380774422, "grad_rms": 8.90960064614394e-05, "grad_to_weight": 0.0021690866055150432, "update_to_weight": 7.652548047644557e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.059032430001572, "grad_norm": 0.02186003738984465, "grad_rms": 0.000497333071074686, "grad_to_weight": 0.004320991749372489, "update_to_weight": 3.177769181653328e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634865493895646, "grad_norm": 0.1031464796382483, "grad_rms": 0.00024854645564396146, "grad_to_weight": 0.0062006199975646985, "update_to_weight": 8.685138070408998e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55474342779186, "grad_norm": 0.01780913388477395, "grad_rms": 2.4570376832001337e-05, "grad_to_weight": 0.0007252828333207909, "update_to_weight": 9.578494759665655e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.821814562573817, "grad_norm": 0.00736661106709727, "grad_rms": 8.942894220208247e-06, "grad_to_weight": 0.00023900640412139677, "update_to_weight": 8.532174979798298e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.591059276002774, "grad_norm": 0.00327747616608542, "grad_rms": 2.9421976878791035e-06, "grad_to_weight": 9.474922811511409e-05, "update_to_weight": 9.593973312895796e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01806859283383, "grad_norm": 0.0027235004178719165, "grad_rms": 2.1906797395715727e-06, "grad_to_weight": 7.357219113260672e-05, "update_to_weight": 0.00010237910619389457, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65751368157445, "grad_norm": 0.0010411574742610413, "grad_rms": 1.881050221392717e-06, "grad_to_weight": 4.2224755005978385e-05, "update_to_weight": 6.775959627154673e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26542682075186, "grad_norm": 0.002891776577265044, "grad_rms": 2.3260346565807234e-06, "grad_to_weight": 7.759944871085471e-05, "update_to_weight": 0.00010322244365504866, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.36892460401037, "grad_norm": 0.007684278014579807, "grad_rms": 3.4115531929109616e-06, "grad_to_weight": 0.00012126255988402081, "update_to_weight": 0.0001129871194943353, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74355091028194, "grad_norm": 0.022202991818160146, "grad_rms": 1.65672287771816e-05, "grad_to_weight": 0.00043755297806053324, "update_to_weight": 8.938083979296489e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.052903857546646, "grad_norm": 0.04128432671437852, "grad_rms": 3.545335041583692e-05, "grad_to_weight": 0.0010307449083145528, "update_to_weight": 0.00010273306457875553, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.194599213407653, "grad_norm": 0.12683399186023953, "grad_rms": 0.00019771621982499508, "grad_to_weight": 0.004498520830185269, "update_to_weight": 8.42237199009247e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40418861952414, "grad_norm": 0.007493760631934307, "grad_rms": 0.0006379112677876886, "grad_to_weight": 0.0008916697341281968, "update_to_weight": 5.58419121397419e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.304531202262597, "grad_norm": 0.11144404900974811, "grad_rms": 0.002579889249981263, "grad_to_weight": 0.0854284273281145, "update_to_weight": 0.00020369789612944018, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41550}
{"step": 41600, "lr": 1.6769474033843667e-05, "window_steps": 50, "elapsed_seconds": 356.268514408, "steps_per_second": 1.684123001991913, "eta_seconds": 4987.759201712, "images_seen": 42598400, "images_per_second": 1724.5419540397188, "effective_batch_size": 1024, "loss": 0.36693546175956726, "edge_loss": 0.19071060419082642, "diag_loss": 0.17622481286525726, "jvp_rms": 2.008125066757202, "grad_norm": 0.17939548194408417, "grad_norm_postclip": 0.17939548194408417, "clip_fraction": 0.0, "ema_updates": 41600, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.36698391288518906, "updates_nfe_3": 12, "loss_nfe_4": 0.3919234871864319, "updates_nfe_4": 8, "loss_nfe_2": 0.3511742323637009, "updates_nfe_2": 10, "loss_nfe_5": 0.31255470402538776, "updates_nfe_5": 8, "loss_nfe_1": 0.3996165618300438, "updates_nfe_1": 12, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47929516664633, "grad_norm": 0.016148766544737948, "grad_rms": 4.702363083503054e-05, "grad_to_weight": 0.0009799427937562909, "update_to_weight": 5.055556924317705e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245350333886403, "grad_norm": 0.010177696304889714, "grad_rms": 2.963645752549051e-05, "grad_to_weight": 0.000626499034844445, "update_to_weight": 5.395538555486549e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.057251835465516, "grad_norm": 0.0027722756308381016, "grad_rms": 4.9631695354505414e-05, "grad_to_weight": 0.0005481782835880697, "update_to_weight": 3.364370776700604e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.83635890684549, "grad_norm": 0.009542819185647373, "grad_rms": 3.053706048153402e-05, "grad_to_weight": 0.000743421031999837, "update_to_weight": 6.17038098189379e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.058764739810404, "grad_norm": 0.019345771481567964, "grad_rms": 0.00044013154102412304, "grad_to_weight": 0.0038242085719710733, "update_to_weight": 4.3435657923529886e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634592142932714, "grad_norm": 0.0822750566968169, "grad_rms": 0.00019825372423390924, "grad_to_weight": 0.004946021879578926, "update_to_weight": 9.79698123089239e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55489256794738, "grad_norm": 0.015082106978614855, "grad_rms": 2.0808033354274927e-05, "grad_to_weight": 0.0006142200352487885, "update_to_weight": 0.00010607302222200952, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.821893772806813, "grad_norm": 0.007128364847801726, "grad_rms": 8.653668860253941e-06, "grad_to_weight": 0.00023127601763688052, "update_to_weight": 9.012666068358592e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59290134178091, "grad_norm": 0.0025651769483672595, "grad_rms": 2.3027650863137423e-06, "grad_to_weight": 7.415327563950435e-05, "update_to_weight": 0.00010029858924259448, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01964645472472, "grad_norm": 0.0020669079627532905, "grad_rms": 1.6625418405849983e-06, "grad_to_weight": 5.5832731014358364e-05, "update_to_weight": 0.00010528844834868385, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.657521107523536, "grad_norm": 0.0007913356155459838, "grad_rms": 1.4296992257345157e-06, "grad_to_weight": 3.209307262052917e-05, "update_to_weight": 6.956099158403963e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26696800515002, "grad_norm": 0.0022462383959352336, "grad_rms": 1.806788393323641e-06, "grad_to_weight": 6.027424596561813e-05, "update_to_weight": 0.00010628126829930539, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.37091448554593, "grad_norm": 0.006242218251165269, "grad_rms": 2.7713286225725455e-06, "grad_to_weight": 9.850289051121452e-05, "update_to_weight": 0.000116552885783695, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74425575557789, "grad_norm": 0.020600253191989208, "grad_rms": 1.5371311681536783e-05, "grad_to_weight": 0.00040596226873866007, "update_to_weight": 9.243874508988423e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05325434460663, "grad_norm": 0.033628602163105564, "grad_rms": 2.887891631930422e-05, "grad_to_weight": 0.0008395972490468512, "update_to_weight": 0.00011384490381115663, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.194402217572375, "grad_norm": 0.10669360381897329, "grad_rms": 0.00016632020893766436, "grad_to_weight": 0.003784212305536157, "update_to_weight": 9.993891187636397e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.404323881913962, "grad_norm": 0.010273221250712998, "grad_rms": 0.0008745146681599, "grad_to_weight": 0.0012223733158143618, "update_to_weight": 4.894675250286863e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3048151376387596, "grad_norm": 0.17387713052279413, "grad_rms": 0.004025192406766444, "grad_to_weight": 0.13325805741144944, "update_to_weight": 0.0001505930930203874, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41600}
{"step": 41650, "lr": 1.6691187665922593e-05, "window_steps": 50, "elapsed_seconds": 383.740873655, "steps_per_second": 1.6938513581026002, "eta_seconds": 4929.594300029615, "images_seen": 42649600, "images_per_second": 1734.5037906970626, "effective_batch_size": 1024, "loss": 0.3370065987110138, "edge_loss": 0.17237140238285065, "diag_loss": 0.16463521122932434, "jvp_rms": 1.577578067779541, "grad_norm": 0.16617143154144287, "grad_norm_postclip": 0.16617143154144287, "clip_fraction": 0.0, "ema_updates": 41650, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.40317683666944504, "updates_nfe_1": 4, "loss_nfe_5": 0.2604319378733635, "updates_nfe_5": 8, "loss_nfe_3": 0.3456557015577952, "updates_nfe_3": 15, "loss_nfe_4": 0.32349875569343567, "updates_nfe_4": 11, "loss_nfe_2": 0.3675704474250476, "updates_nfe_2": 12, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.48144412629795, "grad_norm": 0.010442466732425689, "grad_rms": 3.0407443148819237e-05, "grad_to_weight": 0.0006335893051849496, "update_to_weight": 4.501147505395495e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245873499498842, "grad_norm": 0.004054068958190003, "grad_rms": 1.1805052821932234e-05, "grad_to_weight": 0.00024954453562100337, "update_to_weight": 5.833735515604062e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.057621241580969, "grad_norm": 0.0006249400307903267, "grad_rms": 1.1188221285790353e-05, "grad_to_weight": 0.00012356402366638588, "update_to_weight": 2.7967491771348553e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.837425141083772, "grad_norm": 0.006390087622843559, "grad_rms": 2.0448306566948545e-05, "grad_to_weight": 0.0004977701955506079, "update_to_weight": 5.913500729181934e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.058650684404506, "grad_norm": 0.014378479069779355, "grad_rms": 0.0003271217256233284, "grad_to_weight": 0.0028423546053708117, "update_to_weight": 2.2560283129730883e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634860907507143, "grad_norm": 0.06462536541302374, "grad_rms": 0.000155724224175526, "grad_to_weight": 0.003884935724581801, "update_to_weight": 8.026589193539137e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.555684238439962, "grad_norm": 0.015536357198624992, "grad_rms": 2.143473980467737e-05, "grad_to_weight": 0.0006326990137095861, "update_to_weight": 9.949591201871327e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.822538338509798, "grad_norm": 0.0064744038364026755, "grad_rms": 7.85977542732881e-06, "grad_to_weight": 0.00021005420661002246, "update_to_weight": 8.972991334392064e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59352239931903, "grad_norm": 0.0035890119843321167, "grad_rms": 3.2218640890025425e-06, "grad_to_weight": 0.00010374809315176203, "update_to_weight": 0.00011064089090927516, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.01995476471102, "grad_norm": 0.0030159286446484763, "grad_rms": 2.4258978388508873e-06, "grad_to_weight": 8.146764802434028e-05, "update_to_weight": 0.0001142758093683619, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.657382489438508, "grad_norm": 0.0010663500637290875, "grad_rms": 1.926565455320186e-06, "grad_to_weight": 4.3246685417068784e-05, "update_to_weight": 7.434307164704129e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.267410203571295, "grad_norm": 0.003214430408570846, "grad_rms": 2.5855650779819684e-06, "grad_to_weight": 8.625312011251081e-05, "update_to_weight": 0.00011230414605433533, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.37251135399717, "grad_norm": 0.008844938500915459, "grad_rms": 3.9268462341740585e-06, "grad_to_weight": 0.00013957058528906748, "update_to_weight": 0.00011935048109263683, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74475371205259, "grad_norm": 0.017878146779640624, "grad_rms": 1.3340154796980085e-05, "grad_to_weight": 0.00035231517490633347, "update_to_weight": 8.575565220166198e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.053941595112555, "grad_norm": 0.023818111954076873, "grad_rms": 2.0454054517920056e-05, "grad_to_weight": 0.000594650888415516, "update_to_weight": 9.327420673462417e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19463493228532, "grad_norm": 0.059258341768022256, "grad_rms": 9.237535739142706e-05, "grad_to_weight": 0.002101759498228731, "update_to_weight": 7.0227557043596e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.404480021026174, "grad_norm": 0.004148933342383764, "grad_rms": 0.0003531806603386902, "grad_to_weight": 0.0004936573508419366, "update_to_weight": 5.0049552114320255e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.305030731920875, "grad_norm": 0.06282086264515634, "grad_rms": 0.0014542801491231878, "grad_to_weight": 0.048137458458691086, "update_to_weight": 9.38900430776083e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41650}
{"step": 41700, "lr": 1.6613320237035577e-05, "window_steps": 50, "elapsed_seconds": 411.205829186, "steps_per_second": 1.7023104983353001, "eta_seconds": 4875.726260348286, "images_seen": 42700800, "images_per_second": 1743.1659502953473, "effective_batch_size": 1024, "loss": 0.37033262848854065, "edge_loss": 0.1907452791929245, "diag_loss": 0.17958731949329376, "jvp_rms": 2.1878905296325684, "grad_norm": 0.2087913453578949, "grad_norm_postclip": 0.2087913453578949, "clip_fraction": 0.0, "ema_updates": 41700, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3896505998240577, "updates_nfe_5": 9, "loss_nfe_4": 0.3673128734032313, "updates_nfe_4": 6, "loss_nfe_3": 0.3488965998093287, "updates_nfe_3": 15, "loss_nfe_2": 0.3664595236380895, "updates_nfe_2": 12, "loss_nfe_1": 0.39686712995171547, "updates_nfe_1": 8, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.479778498670253, "grad_norm": 0.006276150022163016, "grad_rms": 1.827553583673728e-05, "grad_to_weight": 0.00038083946472153345, "update_to_weight": 3.6441437762273106e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.246743213304704, "grad_norm": 0.004810910829299766, "grad_rms": 1.400890242548948e-05, "grad_to_weight": 0.00029611539778384867, "update_to_weight": 5.315055545924914e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.057313687885189, "grad_norm": 0.000583474321749739, "grad_rms": 1.0445866010626446e-05, "grad_to_weight": 0.00011537238102264717, "update_to_weight": 3.438217212176227e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.837054881371495, "grad_norm": 0.005762631171828233, "grad_rms": 1.8440443353632943e-05, "grad_to_weight": 0.00044890601661216553, "update_to_weight": 7.164948145182218e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.058218758724102, "grad_norm": 0.01813371947778115, "grad_rms": 0.0004125564031322932, "grad_to_weight": 0.003585001033517033, "update_to_weight": 2.7168473641921183e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634863659340397, "grad_norm": 0.05825734325389793, "grad_rms": 0.00014037954791838286, "grad_to_weight": 0.003502123278370647, "update_to_weight": 7.900250291386284e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55520451651047, "grad_norm": 0.014155863551069977, "grad_rms": 1.9530141335483408e-05, "grad_to_weight": 0.0005764913723912107, "update_to_weight": 0.00010760102045060847, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823063090188146, "grad_norm": 0.0073285216864356255, "grad_rms": 8.896654599429062e-06, "grad_to_weight": 0.00023776097998412434, "update_to_weight": 9.558676543199756e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59404640794424, "grad_norm": 0.004212903112638153, "grad_rms": 3.7819325508833002e-06, "grad_to_weight": 0.00012178116034644313, "update_to_weight": 0.000109856432087369, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.02051697042888, "grad_norm": 0.003559252326322141, "grad_rms": 2.862926661634019e-06, "grad_to_weight": 9.614269647193712e-05, "update_to_weight": 0.00010929908207900249, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65618563838241, "grad_norm": 0.001230713287461097, "grad_rms": 2.223519072840296e-06, "grad_to_weight": 4.9914991130876275e-05, "update_to_weight": 7.480660147789576e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26946553847315, "grad_norm": 0.0036670418145787432, "grad_rms": 2.949628410057849e-06, "grad_to_weight": 9.839265901984206e-05, "update_to_weight": 0.00011004768203658656, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.3741736724975, "grad_norm": 0.00961981035215067, "grad_rms": 4.270862488292223e-06, "grad_to_weight": 0.00015179385851819604, "update_to_weight": 0.00012265842553605362, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74487880180004, "grad_norm": 0.01863479262852709, "grad_rms": 1.3904741992456771e-05, "grad_to_weight": 0.0003672250888865276, "update_to_weight": 9.22762688701283e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05440483492031, "grad_norm": 0.0257425573735589, "grad_rms": 2.210669229217946e-05, "grad_to_weight": 0.0006426897985291238, "update_to_weight": 0.0001049428765842046, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.194936917351622, "grad_norm": 0.05813903400112164, "grad_rms": 9.06305151985218e-05, "grad_to_weight": 0.0020620380946957157, "update_to_weight": 8.303080240278095e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.404521778669206, "grad_norm": 0.007079827113059258, "grad_rms": 0.000602674906663455, "grad_to_weight": 0.0008423831003719876, "update_to_weight": 7.858372948580916e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3050712888984477, "grad_norm": 0.08351415583796538, "grad_rms": 0.0019333223692256604, "grad_to_weight": 0.0639920259899794, "update_to_weight": 0.00029983705929360595, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41700}
{"step": 41750, "lr": 1.6535872598728358e-05, "window_steps": 50, "elapsed_seconds": 438.666152545, "steps_per_second": 1.709728447587627, "eta_seconds": 4825.327677994999, "images_seen": 42752000, "images_per_second": 1750.7619303297301, "effective_batch_size": 1024, "loss": 0.38218387961387634, "edge_loss": 0.19776879251003265, "diag_loss": 0.1844150274991989, "jvp_rms": 2.5868358612060547, "grad_norm": 0.28679612278938293, "grad_norm_postclip": 0.27899837493896484, "clip_fraction": 0.019999999552965164, "ema_updates": 41750, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.37211285396055743, "updates_nfe_2": 11, "loss_nfe_1": 0.4026401311159134, "updates_nfe_1": 10, "loss_nfe_5": 0.4460581053387035, "updates_nfe_5": 11, "loss_nfe_4": 0.2684016498652371, "updates_nfe_4": 11, "loss_nfe_3": 0.4472133346966335, "updates_nfe_3": 7, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47936183397519, "grad_norm": 0.013063055242753887, "grad_rms": 3.803834092288793e-05, "grad_to_weight": 0.0007926918150326691, "update_to_weight": 5.3910644587916724e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.247443834411115, "grad_norm": 0.007586572635729658, "grad_rms": 2.2091358490902327e-05, "grad_to_weight": 0.000466939459095821, "update_to_weight": 5.688737427760442e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.057605779487279, "grad_norm": 0.000662032736498345, "grad_rms": 1.1852287242687334e-05, "grad_to_weight": 0.00013089844589774639, "update_to_weight": 2.838187946910223e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.837640873359847, "grad_norm": 0.005867239397995004, "grad_rms": 1.877519010584273e-05, "grad_to_weight": 0.000457034080940094, "update_to_weight": 8.71274826734215e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.058118266142767, "grad_norm": 0.020287663263721962, "grad_rms": 0.00046156032105248084, "grad_to_weight": 0.004010911211689201, "update_to_weight": 4.8119518392567165e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.634808622588846, "grad_norm": 0.05803207562739508, "grad_rms": 0.0001398367327846529, "grad_to_weight": 0.003488592922469321, "update_to_weight": 0.00011184441128105268, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.555361110709484, "grad_norm": 0.017230657760575377, "grad_rms": 2.3772282076139775e-05, "grad_to_weight": 0.0007017065512858804, "update_to_weight": 0.00010641129682909502, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823383877380234, "grad_norm": 0.008911074840670621, "grad_rms": 1.0817837260937096e-05, "grad_to_weight": 0.00028910112128246963, "update_to_weight": 8.55271067126412e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59529376354445, "grad_norm": 0.005687007757913274, "grad_rms": 5.105239589359952e-06, "grad_to_weight": 0.000164386745688111, "update_to_weight": 9.048414704888622e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.02158695127243, "grad_norm": 0.004770358625293087, "grad_rms": 3.837094322565121e-06, "grad_to_weight": 0.00012885343439144847, "update_to_weight": 9.326719729603609e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65636882100647, "grad_norm": 0.001647422005839631, "grad_rms": 2.9763831172717617e-06, "grad_to_weight": 6.681527267048658e-05, "update_to_weight": 5.9158427478483145e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.27043338999595, "grad_norm": 0.004845640504906953, "grad_rms": 3.897648191896203e-06, "grad_to_weight": 0.00013001299057090142, "update_to_weight": 9.970629477117407e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.37546420256521, "grad_norm": 0.01280499665923248, "grad_rms": 5.684974848012126e-06, "grad_to_weight": 0.00020204974938415016, "update_to_weight": 0.0001107962950658563, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.744907668621046, "grad_norm": 0.02251285792395339, "grad_rms": 1.6798441881569276e-05, "grad_to_weight": 0.00044364762807272954, "update_to_weight": 9.890683534242558e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05503873334399, "grad_norm": 0.02530957775776, "grad_rms": 2.1734866486516353e-05, "grad_to_weight": 0.0006318700108181629, "update_to_weight": 0.00014164375195858448, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.195024589829565, "grad_norm": 0.049182812403599424, "grad_rms": 7.666903490285838e-05, "grad_to_weight": 0.0017443791278458584, "update_to_weight": 0.00012069810401750638, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.404771412425038, "grad_norm": 0.006000927718986498, "grad_rms": 0.0005108328911398456, "grad_to_weight": 0.0007139905923098797, "update_to_weight": 1.0412251670612293e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3053678919916363, "grad_norm": 0.06297261020977457, "grad_rms": 0.0014577930501183335, "grad_to_weight": 0.04824127404703934, "update_to_weight": 0.00031306439271342184, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41750}
{"step": 41800, "lr": 1.6458845597955905e-05, "window_steps": 50, "elapsed_seconds": 466.13313701299995, "steps_per_second": 1.7162478624163742, "eta_seconds": 4777.864654383249, "images_seen": 42803200, "images_per_second": 1757.4378111143672, "effective_batch_size": 1024, "loss": 0.37212374806404114, "edge_loss": 0.19583599269390106, "diag_loss": 0.17628775537014008, "jvp_rms": 2.231250047683716, "grad_norm": 0.3096234202384949, "grad_norm_postclip": 0.3096234202384949, "clip_fraction": 0.0, "ema_updates": 41800, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.311943631619215, "updates_nfe_4": 8, "loss_nfe_2": 0.3757927119731903, "updates_nfe_2": 14, "loss_nfe_1": 0.4184451748927434, "updates_nfe_1": 12, "loss_nfe_3": 0.4368758499622345, "updates_nfe_3": 6, "loss_nfe_5": 0.3206943362951279, "updates_nfe_5": 10, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47877292990563, "grad_norm": 0.07620587568224783, "grad_rms": 0.00022190406651893325, "grad_to_weight": 0.004624487272589916, "update_to_weight": 0.00012231287733163696, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248235518686425, "grad_norm": 0.021560131695338023, "grad_rms": 6.278099759430688e-05, "grad_to_weight": 0.001326921416823544, "update_to_weight": 0.00010450873449860961, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0569974405781215, "grad_norm": 0.0021137368870113643, "grad_rms": 3.7841960614866155e-05, "grad_to_weight": 0.00041798258983688936, "update_to_weight": 4.9652467888798416e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.839086126682046, "grad_norm": 0.02352194177660656, "grad_rms": 7.527031003119949e-05, "grad_to_weight": 0.001832057324370114, "update_to_weight": 0.00013616774351644532, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.057751536127125, "grad_norm": 0.05358759850323226, "grad_rms": 0.0012191600800971524, "grad_to_weight": 0.010595142549107093, "update_to_weight": 6.302119038911095e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.635735967538096, "grad_norm": 0.0894235445446964, "grad_rms": 0.00021547904616477564, "grad_to_weight": 0.0053753885442274245, "update_to_weight": 0.00015497465031323102, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.556192534610798, "grad_norm": 0.024559348069543414, "grad_rms": 3.388331182870579e-05, "grad_to_weight": 0.0010001285026140828, "update_to_weight": 0.00016888827361305833, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82350763697855, "grad_norm": 0.009631670573542053, "grad_rms": 1.1692623693383016e-05, "grad_to_weight": 0.0003124780828638421, "update_to_weight": 0.00013770079766182818, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59330714768483, "grad_norm": 0.005807849339012534, "grad_rms": 5.213719347104324e-06, "grad_to_weight": 0.0001678893930036182, "update_to_weight": 0.00015398217931846623, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.021919974438795, "grad_norm": 0.00486981946804366, "grad_rms": 3.917096826572316e-06, "grad_to_weight": 0.0001315388146105321, "update_to_weight": 0.00015126984436610963, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65703222776922, "grad_norm": 0.0016253113973536007, "grad_rms": 2.936436071780594e-06, "grad_to_weight": 6.591674871248876e-05, "update_to_weight": 0.00010154185381026351, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.26997157565324, "grad_norm": 0.004883451904214745, "grad_rms": 3.928062237675257e-06, "grad_to_weight": 0.00013102912875321, "update_to_weight": 0.00014786612501915843, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.37437014294994, "grad_norm": 0.01293585104108093, "grad_rms": 5.743069659698333e-06, "grad_to_weight": 0.00020411802140048527, "update_to_weight": 0.00015878871618013645, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74515063038044, "grad_norm": 0.028990499572026658, "grad_rms": 2.163187027712693e-05, "grad_to_weight": 0.00057129596053796, "update_to_weight": 0.00012783019673962072, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.056036798270185, "grad_norm": 0.045621947184746686, "grad_rms": 3.91783276831375e-05, "grad_to_weight": 0.0011389530974945795, "update_to_weight": 0.00017752719723810491, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.195330900174103, "grad_norm": 0.16170175850103122, "grad_rms": 0.00025207012695072783, "grad_to_weight": 0.00573505446960485, "update_to_weight": 0.00015988261159958345, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.404349754002249, "grad_norm": 0.02659374933853552, "grad_rms": 0.002263810280178951, "grad_to_weight": 0.0031642839859051877, "update_to_weight": 1.2994868901448291e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3049120224422763, "grad_norm": 0.3291018002110197, "grad_rms": 0.007618587121144748, "grad_to_weight": 0.25220228992531774, "update_to_weight": 0.00032811617834104255, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41800}
{"step": 41850, "lr": 1.638224007707315e-05, "window_steps": 50, "elapsed_seconds": 493.59488984399997, "steps_per_second": 1.722059967575113, "eta_seconds": 4732.703943798352, "images_seen": 42854400, "images_per_second": 1763.3894067969156, "effective_batch_size": 1024, "loss": 0.36485472321510315, "edge_loss": 0.18873530626296997, "diag_loss": 0.17611949145793915, "jvp_rms": 2.1227734088897705, "grad_norm": 0.29050323367118835, "grad_norm_postclip": 0.29050323367118835, "clip_fraction": 0.0, "ema_updates": 41850, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3457946976025899, "updates_nfe_5": 9, "loss_nfe_3": 0.3318438885303644, "updates_nfe_3": 13, "loss_nfe_2": 0.3753862449756035, "updates_nfe_2": 13, "loss_nfe_4": 0.38296334279908073, "updates_nfe_4": 9, "loss_nfe_1": 0.4149876882632573, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.476664369647153, "grad_norm": 0.02074669182221857, "grad_rms": 6.0412340137150965e-05, "grad_to_weight": 0.0012591560619780262, "update_to_weight": 7.829296997093119e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25010422893256, "grad_norm": 0.012476204037518118, "grad_rms": 3.632948753438616e-05, "grad_to_weight": 0.0007677614778190046, "update_to_weight": 7.563212978286112e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.05848025632819, "grad_norm": 0.0025072333565840602, "grad_rms": 4.488667748344147e-05, "grad_to_weight": 0.0004956495290156556, "update_to_weight": 4.402614776198022e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.83929767149534, "grad_norm": 0.01028021456044422, "grad_rms": 3.289672870126119e-05, "grad_to_weight": 0.000800683559449473, "update_to_weight": 0.00011014929651779619, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0581684184122855, "grad_norm": 0.026792869435321092, "grad_rms": 0.0006095588859953922, "grad_to_weight": 0.0052969508365502665, "update_to_weight": 3.185223257354609e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63510031529717, "grad_norm": 0.0920375651312311, "grad_rms": 0.00022177790923835964, "grad_to_weight": 0.0055327327991281155, "update_to_weight": 0.00011091405724419775, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.555819702262028, "grad_norm": 0.034290104715905156, "grad_rms": 4.7308353114178965e-05, "grad_to_weight": 0.001396414582435887, "update_to_weight": 0.00014542445771423398, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824092765630894, "grad_norm": 0.021638610572652355, "grad_rms": 2.626876913426637e-05, "grad_to_weight": 0.0007020031615262843, "update_to_weight": 0.00012704692145716447, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.596676917168466, "grad_norm": 0.01342229245191553, "grad_rms": 1.204922196741118e-05, "grad_to_weight": 0.0003879647887585056, "update_to_weight": 0.00013908728475962918, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.02450162666778, "grad_norm": 0.011189518911063767, "grad_rms": 9.000421741507941e-06, "grad_to_weight": 0.0003022192985578029, "update_to_weight": 0.00014395742430133176, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.657778539042578, "grad_norm": 0.004532369671724757, "grad_rms": 8.188593162127106e-06, "grad_to_weight": 0.0001838109489282785, "update_to_weight": 9.556393883437863e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.272514758804675, "grad_norm": 0.01009739989950479, "grad_rms": 8.121962911054515e-06, "grad_to_weight": 0.0002709073955660461, "update_to_weight": 0.00014530346404244103, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.37797391937518, "grad_norm": 0.02057417919325516, "grad_rms": 9.134222705776224e-06, "grad_to_weight": 0.00032462664741268195, "update_to_weight": 0.00015176048105421529, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74572796058797, "grad_norm": 0.031460342828237654, "grad_rms": 2.3474795708282514e-05, "grad_to_weight": 0.0006199604201691924, "update_to_weight": 0.00011473984863595326, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.056693526663715, "grad_norm": 0.03056661861699343, "grad_rms": 2.624940569705561e-05, "grad_to_weight": 0.0007630839174643004, "update_to_weight": 0.00012414529467925212, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.194965059164375, "grad_norm": 0.050650326977928366, "grad_rms": 7.89566821645966e-05, "grad_to_weight": 0.0017964316278329699, "update_to_weight": 8.917393360250349e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.404641603797348, "grad_norm": 0.0027623640402573042, "grad_rms": 0.0002351480429602153, "grad_to_weight": 0.0003286712474460809, "update_to_weight": 5.922699547754785e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3055808383892409, "grad_norm": 0.032039689342065976, "grad_rms": 0.0007417071691203989, "grad_to_weight": 0.024540563402872022, "update_to_weight": 0.0001313202200064605, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41850}
{"step": 41900, "lr": 1.6306056873825808e-05, "window_steps": 50, "elapsed_seconds": 521.044603679, "steps_per_second": 1.7272993399130625, "eta_seconds": 4689.401433111, "images_seen": 42905600, "images_per_second": 1768.754524070976, "effective_batch_size": 1024, "loss": 0.37012362480163574, "edge_loss": 0.19491830468177795, "diag_loss": 0.17520536482334137, "jvp_rms": 2.1370701789855957, "grad_norm": 0.24266213178634644, "grad_norm_postclip": 0.24266213178634644, "clip_fraction": 0.0, "ema_updates": 41900, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.39293273199688306, "updates_nfe_3": 11, "loss_nfe_5": 0.32115156813101337, "updates_nfe_5": 11, "loss_nfe_4": 0.32309434314568836, "updates_nfe_4": 6, "loss_nfe_1": 0.41466279540743145, "updates_nfe_1": 14, "loss_nfe_2": 0.3634261153638363, "updates_nfe_2": 8, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.484332419968702, "grad_norm": 0.10637737433426223, "grad_rms": 0.0003097605235691624, "grad_to_weight": 0.006453241273234661, "update_to_weight": 9.473422213022082e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2429024088095, "grad_norm": 0.054553636560856614, "grad_rms": 0.00015885486109662344, "grad_to_weight": 0.0033586138233071515, "update_to_weight": 0.0001263713209173939, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0587775591002275, "grad_norm": 0.007607034341590741, "grad_rms": 0.00013618776098354735, "grad_to_weight": 0.0015037297554043383, "update_to_weight": 5.875368493917593e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.83875870073061, "grad_norm": 0.025928879791577174, "grad_rms": 8.297252153794249e-05, "grad_to_weight": 0.0020195784028639507, "update_to_weight": 0.00011972550484492167, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.058914987387055, "grad_norm": 0.10478386322674214, "grad_rms": 0.0023839154329093246, "grad_to_weight": 0.020712714779352978, "update_to_weight": 5.4207729913685176e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63631197774023, "grad_norm": 0.2597527321505666, "grad_rms": 0.0006259120150904164, "grad_to_weight": 0.015613600688549346, "update_to_weight": 0.00013718996869483555, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.555640740723913, "grad_norm": 0.0810402202334111, "grad_rms": 0.00011180716381640949, "grad_to_weight": 0.0033002690130993504, "update_to_weight": 0.0001525435597546434, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823676939304203, "grad_norm": 0.040262365731725584, "grad_rms": 4.887757403161475e-05, "grad_to_weight": 0.001306215537199127, "update_to_weight": 0.00012006670077035433, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59605062358997, "grad_norm": 0.026504815275082954, "grad_rms": 2.3793431978836664e-05, "grad_to_weight": 0.0007661225717195057, "update_to_weight": 0.0001427226785986893, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.024478547560264, "grad_norm": 0.020680999769932113, "grad_rms": 1.663500651322658e-05, "grad_to_weight": 0.0005585763954343361, "update_to_weight": 0.00015094266552634453, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65865601378652, "grad_norm": 0.007700952405832298, "grad_rms": 1.391324423637926e-05, "grad_to_weight": 0.00031230219528293585, "update_to_weight": 9.960130410801882e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.27231006600275, "grad_norm": 0.020050248416505274, "grad_rms": 1.6127654209701243e-05, "grad_to_weight": 0.0005379395154472526, "update_to_weight": 0.00015294099639215144, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.37792961978661, "grad_norm": 0.039457288666355554, "grad_rms": 1.7517669048140964e-05, "grad_to_weight": 0.0006225714361934123, "update_to_weight": 0.00016444212568974954, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74575201587075, "grad_norm": 0.058997345755061985, "grad_rms": 4.402210892908349e-05, "grad_to_weight": 0.0011626065909243113, "update_to_weight": 0.00013743173019297125, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.056943415846305, "grad_norm": 0.06925170442806725, "grad_rms": 5.947063060924605e-05, "grad_to_weight": 0.0017288314714664838, "update_to_weight": 0.00017494074064910642, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.195775747787668, "grad_norm": 0.20856456838216025, "grad_rms": 0.0003251226066856881, "grad_to_weight": 0.007397014724750919, "update_to_weight": 0.0001305965233118452, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40528472690438, "grad_norm": 0.026743221873175106, "grad_rms": 0.0022765342273070523, "grad_to_weight": 0.0031817151639816593, "update_to_weight": 1.047350798738947e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3060280771056079, "grad_norm": 0.2982392836193917, "grad_rms": 0.00690413107356212, "grad_to_weight": 0.22835595103004477, "update_to_weight": 0.00027947931068743136, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41900}
{"step": 41950, "lr": 1.6230296821341173e-05, "window_steps": 50, "elapsed_seconds": 548.5057761310001, "steps_per_second": 1.7319781146171755, "eta_seconds": 4647.864734583737, "images_seen": 42956800, "images_per_second": 1773.5455893679878, "effective_batch_size": 1024, "loss": 0.354535847902298, "edge_loss": 0.1845802664756775, "diag_loss": 0.16995558142662048, "jvp_rms": 1.8132030963897705, "grad_norm": 0.290225088596344, "grad_norm_postclip": 0.290225088596344, "clip_fraction": 0.0, "ema_updates": 41950, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.30533476008309257, "updates_nfe_4": 9, "loss_nfe_5": 0.32088732854886487, "updates_nfe_5": 11, "loss_nfe_3": 0.35622692650014703, "updates_nfe_3": 11, "loss_nfe_2": 0.37505110204219816, "updates_nfe_2": 10, "loss_nfe_1": 0.4200013213687473, "updates_nfe_1": 9, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.47931553724765, "grad_norm": 0.045188192129770245, "grad_rms": 0.0001315836016903191, "grad_to_weight": 0.002742115837738095, "update_to_weight": 9.413808884313006e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24893701451152, "grad_norm": 0.03319217873101221, "grad_rms": 9.665238239301072e-05, "grad_to_weight": 0.0020427292383107345, "update_to_weight": 0.00010532658345000577, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.059388511408765, "grad_norm": 0.002804980063167786, "grad_rms": 5.021719861546299e-05, "grad_to_weight": 0.000554410885197419, "update_to_weight": 4.863634793543545e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.839420674882554, "grad_norm": 0.016906842427673166, "grad_rms": 5.4101965019113673e-05, "grad_to_weight": 0.0013167916883312042, "update_to_weight": 0.00010385098684527595, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.05834375890488, "grad_norm": 0.041551547196011424, "grad_rms": 0.0009453304313421386, "grad_to_weight": 0.008214456979690768, "update_to_weight": 3.077936955703222e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.635888226951074, "grad_norm": 0.07240943839458239, "grad_rms": 0.00017448108099531908, "grad_to_weight": 0.004352604285791908, "update_to_weight": 0.00010228313088647273, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.555757562986884, "grad_norm": 0.02571732079763231, "grad_rms": 3.548090924553749e-05, "grad_to_weight": 0.001047303091002831, "update_to_weight": 0.00012776625491970237, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82362743579641, "grad_norm": 0.013178745013910585, "grad_rms": 1.5998689430055537e-05, "grad_to_weight": 0.0004275533449578913, "update_to_weight": 0.00010259359236788508, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59795769634116, "grad_norm": 0.009286413731404276, "grad_rms": 8.33643438568771e-06, "grad_to_weight": 0.0002684093036042513, "update_to_weight": 0.00012212792329878384, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.026159987737046, "grad_norm": 0.007687090392828696, "grad_rms": 6.183201981288285e-06, "grad_to_weight": 0.00020761241228835607, "update_to_weight": 0.0001282819169567205, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.659236441839216, "grad_norm": 0.003089834877683296, "grad_rms": 5.5823780018082755e-06, "grad_to_weight": 0.00012530132005388403, "update_to_weight": 8.74624565707706e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.274684433397624, "grad_norm": 0.007777594282817791, "grad_rms": 6.255999854514405e-06, "grad_to_weight": 0.00020865620731718844, "update_to_weight": 0.00013143723425149882, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.38055287248507, "grad_norm": 0.017839115159997886, "grad_rms": 7.919949039756995e-06, "grad_to_weight": 0.0002814603904747926, "update_to_weight": 0.0001412914145626178, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7462836347094, "grad_norm": 0.026640626545176082, "grad_rms": 1.987846315289786e-05, "grad_to_weight": 0.0005249768975585525, "update_to_weight": 0.0001142481452205434, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.057837823499035, "grad_norm": 0.03198264937283506, "grad_rms": 2.7465437023756914e-05, "grad_to_weight": 0.0007984117743387826, "update_to_weight": 0.00012557334188676697, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.195285440898626, "grad_norm": 0.0725351282839023, "grad_rms": 0.00011307198613300162, "grad_to_weight": 0.0025725977641172086, "update_to_weight": 9.981147124514304e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40494842095919, "grad_norm": 0.007702365479779384, "grad_rms": 0.0006556688916952943, "grad_to_weight": 0.0009164084172809681, "update_to_weight": 6.128310146934836e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.306017580299582, "grad_norm": 0.0732178902060661, "grad_rms": 0.0016949675602006852, "grad_to_weight": 0.056061948407517576, "update_to_weight": 0.0001589371708942852, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 41950}
{"step": 42000, "lr": 1.6154960748119043e-05, "window_steps": 50, "elapsed_seconds": 575.9582640550001, "steps_per_second": 1.7362369157785134, "eta_seconds": 4607.666112440001, "images_seen": 43008000, "images_per_second": 1777.9066017571977, "effective_batch_size": 1024, "loss": 0.3711022138595581, "edge_loss": 0.19140301644802094, "diag_loss": 0.17969918251037598, "jvp_rms": 2.2016797065734863, "grad_norm": 0.2867734134197235, "grad_norm_postclip": 0.27830836176872253, "clip_fraction": 0.03999999910593033, "ema_updates": 42000, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.34418162405490876, "updates_nfe_5": 10, "loss_nfe_1": 0.4146569456372942, "updates_nfe_1": 7, "loss_nfe_3": 0.370359280705452, "updates_nfe_3": 10, "loss_nfe_2": 0.37312300909649243, "updates_nfe_2": 11, "loss_nfe_4": 0.3668959215283394, "updates_nfe_4": 12, "gpu_allocated_gb": 0.8970851898193359, "gpu_reserved_gb": 43.7734375, "gpu_peak_allocated_gb": 37.95361328125, "gradient_groups": {"t_emb": {"weight_norm": 16.484599931818465, "grad_norm": 0.03709686532733429, "grad_rms": 0.00010802244836822116, "grad_to_weight": 0.0022503952465191567, "update_to_weight": 6.687291868707075e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.244043755159343, "grad_norm": 0.014514829357110206, "grad_rms": 4.226576534073549e-05, "grad_to_weight": 0.0008935477874774922, "update_to_weight": 8.991928850907263e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058728732513749, "grad_norm": 0.0037607365231662007, "grad_rms": 6.732798403956647e-05, "grad_to_weight": 0.0007434153365439386, "update_to_weight": 4.959980932088199e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.839643504104776, "grad_norm": 0.01052355370633499, "grad_rms": 3.367541496483071e-05, "grad_to_weight": 0.0008196141663100426, "update_to_weight": 0.00010535423291851516, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.05833961113546, "grad_norm": 0.020158595805283917, "grad_rms": 0.00045862393469887756, "grad_to_weight": 0.003985219924915018, "update_to_weight": 2.615675943573195e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.6367494753473, "grad_norm": 0.04478515416630683, "grad_rms": 0.00010791634743660547, "grad_to_weight": 0.0026919413694766786, "update_to_weight": 9.481231997857073e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.556093113204614, "grad_norm": 0.02839899539891425, "grad_rms": 3.9180682402425096e-05, "grad_to_weight": 0.0011564948572231623, "update_to_weight": 0.0001303047970704615, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823315561869112, "grad_norm": 0.01543692934930398, "grad_rms": 1.8740072605740403e-05, "grad_to_weight": 0.00050081988481475, "update_to_weight": 0.00010992832958205176, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59565719957221, "grad_norm": 0.010952771450682703, "grad_rms": 9.83232743887695e-06, "grad_to_weight": 0.0003165938252740618, "update_to_weight": 0.00013480249511487526, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.025312683309984, "grad_norm": 0.00902717500349521, "grad_rms": 7.261114871124622e-06, "grad_to_weight": 0.00024381090527736225, "update_to_weight": 0.00013388102017872718, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.659522319291284, "grad_norm": 0.003587619348035933, "grad_rms": 6.481720907478968e-06, "grad_to_weight": 0.00014548616561113668, "update_to_weight": 8.887508058047539e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.27522806037711, "grad_norm": 0.009490856019750442, "grad_rms": 7.634082174991607e-06, "grad_to_weight": 0.00025461563922231367, "update_to_weight": 0.00013446856815992482, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.37971506128755, "grad_norm": 0.022134209493603063, "grad_rms": 9.826822107059177e-06, "grad_to_weight": 0.0003492317608591251, "update_to_weight": 0.0001470338119706484, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7469547654709, "grad_norm": 0.03126031921299338, "grad_rms": 2.3325543885747482e-05, "grad_to_weight": 0.0006160038441215676, "update_to_weight": 0.00010530305781139793, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05872002204623, "grad_norm": 0.032545385665097966, "grad_rms": 2.7948692742067005e-05, "grad_to_weight": 0.000812441976358373, "update_to_weight": 0.00011801437238513045, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.196019274530357, "grad_norm": 0.056100977333091766, "grad_rms": 8.745347366350479e-05, "grad_to_weight": 0.0019896772231166736, "update_to_weight": 8.000767688785228e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.405079586417008, "grad_norm": 0.00640638088401179, "grad_rms": 0.0005453473566043041, "grad_to_weight": 0.0007622034768552094, "update_to_weight": 5.00065977063346e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3060274838101253, "grad_norm": 0.06542019988776933, "grad_rms": 0.0015144538620210978, "grad_to_weight": 0.05009098253959894, "update_to_weight": 0.00011292153515970481, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42000}
{"step": 42050, "lr": 1.6080049478022643e-05, "window_steps": 50, "elapsed_seconds": 626.310831899, "steps_per_second": 1.6764838583684671, "eta_seconds": 4742.067727235286, "images_seen": 43059200, "images_per_second": 1716.7194709693104, "effective_batch_size": 1024, "loss": 0.35837888717651367, "edge_loss": 0.18472467362880707, "diag_loss": 0.17365427315235138, "jvp_rms": 1.9286718368530273, "grad_norm": 0.22062908113002777, "grad_norm_postclip": 0.21884885430335999, "clip_fraction": 0.019999999552965164, "ema_updates": 42050, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3640437126159668, "updates_nfe_4": 13, "loss_nfe_1": 0.401017002761364, "updates_nfe_1": 8, "loss_nfe_3": 0.3312023530403773, "updates_nfe_3": 12, "loss_nfe_5": 0.32428772416379714, "updates_nfe_5": 9, "loss_nfe_2": 0.38565298169851303, "updates_nfe_2": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.481172859899583, "grad_norm": 0.14067009361789115, "grad_rms": 0.0004096175725552444, "grad_to_weight": 0.008535199212682016, "update_to_weight": 1.9821467713817124e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.246058528905078, "grad_norm": 0.06034325963840447, "grad_rms": 0.00017571367799253652, "grad_to_weight": 0.003714332281337126, "update_to_weight": 2.8958786043482033e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0589805896972555, "grad_norm": 0.004503910900012864, "grad_rms": 8.063293967118851e-05, "grad_to_weight": 0.0008902803282513466, "update_to_weight": 1.8082541684698626e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.838995208879581, "grad_norm": 0.04836494102553896, "grad_rms": 0.00015476800938490348, "grad_to_weight": 0.003767034743660413, "update_to_weight": 3.472591704088048e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.058075466632291, "grad_norm": 0.029284170495562707, "grad_rms": 0.0006662379476623764, "grad_to_weight": 0.005789587499978595, "update_to_weight": 1.803400882724412e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63615238455423, "grad_norm": 0.1357078646072095, "grad_rms": 0.0003270078520316754, "grad_to_weight": 0.0081574069213989, "update_to_weight": 5.4360799511436746e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55626212935562, "grad_norm": 0.019488580854786774, "grad_rms": 2.6887426340953452e-05, "grad_to_weight": 0.0007936297776968784, "update_to_weight": 6.693761296355551e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823963068021445, "grad_norm": 0.008062439605223086, "grad_rms": 9.787613855218557e-06, "grad_to_weight": 0.00026156401717167656, "update_to_weight": 5.727897774644031e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.596892147836954, "grad_norm": 0.002556583579570961, "grad_rms": 2.295050799917036e-06, "grad_to_weight": 7.389633637167037e-05, "update_to_weight": 6.586433104801235e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.02686221367705, "grad_norm": 0.0021559475891381307, "grad_rms": 1.734161916080601e-06, "grad_to_weight": 5.822658092647567e-05, "update_to_weight": 7.270702198381893e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65989977145652, "grad_norm": 0.0008278271766978722, "grad_rms": 1.495628214269554e-06, "grad_to_weight": 3.356977053313373e-05, "update_to_weight": 4.463848926000542e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.27694894772858, "grad_norm": 0.0021786044260048583, "grad_rms": 1.7523862105073671e-06, "grad_to_weight": 5.8443743050424964e-05, "update_to_weight": 7.960245367804904e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.38244994426691, "grad_norm": 0.006465600155318949, "grad_rms": 2.8705024482602787e-06, "grad_to_weight": 0.00010200931268835842, "update_to_weight": 9.394716698372986e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74783516010806, "grad_norm": 0.02130159334084181, "grad_rms": 1.5894631367091957e-05, "grad_to_weight": 0.0004197537347876191, "update_to_weight": 7.63745544372474e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.059003418492736, "grad_norm": 0.0511284226011346, "grad_rms": 4.390707144693975e-05, "grad_to_weight": 0.0012763278723387263, "update_to_weight": 8.638005874317037e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.195585254291693, "grad_norm": 0.16916877653265705, "grad_rms": 0.00026371014992031906, "grad_to_weight": 0.0059998320661532505, "update_to_weight": 6.16792393604914e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.405102733050025, "grad_norm": 0.005079983092766265, "grad_rms": 0.0004324368783861309, "grad_to_weight": 0.0006043927426123025, "update_to_weight": 3.808017526259854e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.306236033387785, "grad_norm": 0.10174651623222723, "grad_rms": 0.002355394889031773, "grad_to_weight": 0.0778929026849327, "update_to_weight": 8.619934928449158e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42050}
{"step": 42100, "lr": 1.6005563830269627e-05, "window_steps": 50, "elapsed_seconds": 653.7662803000001, "steps_per_second": 1.6825584817486035, "eta_seconds": 4695.230558518182, "images_seen": 43110400, "images_per_second": 1722.93988531057, "effective_batch_size": 1024, "loss": 0.37827110290527344, "edge_loss": 0.19733552634716034, "diag_loss": 0.1809355467557907, "jvp_rms": 2.3573436737060547, "grad_norm": 0.17480894923210144, "grad_norm_postclip": 0.17480894923210144, "clip_fraction": 0.0, "ema_updates": 42100, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3732543006539345, "updates_nfe_5": 10, "loss_nfe_1": 0.39575248105185373, "updates_nfe_1": 14, "loss_nfe_2": 0.36969170891321623, "updates_nfe_2": 13, "loss_nfe_4": 0.3663649757703145, "updates_nfe_4": 9, "loss_nfe_3": 0.3843000940978527, "updates_nfe_3": 4, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.482836495147954, "grad_norm": 0.020797635733826203, "grad_rms": 6.0560683831766484e-05, "grad_to_weight": 0.0012617752860648952, "update_to_weight": 6.764831447008123e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.243871853801803, "grad_norm": 0.01238919181697154, "grad_rms": 3.607611644714038e-05, "grad_to_weight": 0.0007626994308054645, "update_to_weight": 7.69952380855764e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058642766413101, "grad_norm": 0.0033413677858380154, "grad_rms": 5.982007927687023e-05, "grad_to_weight": 0.0006605265364897979, "update_to_weight": 4.168968921390714e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.839465835250714, "grad_norm": 0.013279611560090703, "grad_rms": 4.2494811385683634e-05, "grad_to_weight": 0.0010342806881912152, "update_to_weight": 9.054766160458016e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.057837894483187, "grad_norm": 0.06159652614687767, "grad_rms": 0.001401369493100239, "grad_to_weight": 0.012178430276317037, "update_to_weight": 3.2799989685998575e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636485327224882, "grad_norm": 0.13022717736897904, "grad_rms": 0.00031380133841790253, "grad_to_weight": 0.007827805862087225, "update_to_weight": 8.950607812466298e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.557146959630206, "grad_norm": 0.04664347762500879, "grad_rms": 6.435168770230375e-05, "grad_to_weight": 0.0018993850426389748, "update_to_weight": 0.00013133289281378843, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824119497131424, "grad_norm": 0.034061098433562836, "grad_rms": 4.134937999862118e-05, "grad_to_weight": 0.0011050144818161847, "update_to_weight": 0.00011958525317886561, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59723968961433, "grad_norm": 0.02423573934411825, "grad_rms": 2.175647743839212e-05, "grad_to_weight": 0.0007005107795172898, "update_to_weight": 0.00014803987194538312, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0273550826151, "grad_norm": 0.020496137309941438, "grad_rms": 1.6486310209362827e-05, "grad_to_weight": 0.0005535404098999413, "update_to_weight": 0.00015807936506870762, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.659643599633338, "grad_norm": 0.007942119546681779, "grad_rms": 1.4348958827977943e-05, "grad_to_weight": 0.0003220695187500548, "update_to_weight": 0.00010777680560901437, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.277224020359434, "grad_norm": 0.01962951957517816, "grad_rms": 1.578923599522235e-05, "grad_to_weight": 0.0005265821179296303, "update_to_weight": 0.00016436494431908773, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.382935277367835, "grad_norm": 0.03427849180793176, "grad_rms": 1.5218462678424645e-05, "grad_to_weight": 0.0005408157835847591, "update_to_weight": 0.00017184286685462848, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74837637873871, "grad_norm": 0.040435481520134105, "grad_rms": 3.0171783989562905e-05, "grad_to_weight": 0.0007967837476880296, "update_to_weight": 0.00011926860433292414, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05937365921677, "grad_norm": 0.04101807865115166, "grad_rms": 3.522470708713678e-05, "grad_to_weight": 0.001023932101387569, "update_to_weight": 0.0001257553367172977, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19594459321954, "grad_norm": 0.10092137708421516, "grad_rms": 0.00015732212543315138, "grad_to_weight": 0.003579286969817793, "update_to_weight": 9.254946433145458e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.405311957566115, "grad_norm": 0.011177496527378063, "grad_rms": 0.0009514916916463811, "grad_to_weight": 0.0013298134065466234, "update_to_weight": 6.562886519692731e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3063513374262488, "grad_norm": 0.13082959419827855, "grad_rms": 0.0030286575788539966, "grad_to_weight": 0.10014885770013203, "update_to_weight": 0.0002119390930798575, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42100}
{"step": 42150, "lr": 1.593150461942311e-05, "window_steps": 50, "elapsed_seconds": 681.220526854, "steps_per_second": 1.6881464293375137, "eta_seconds": 4650.070552872957, "images_seen": 43161600, "images_per_second": 1728.661943641614, "effective_batch_size": 1024, "loss": 0.3639601469039917, "edge_loss": 0.18647685647010803, "diag_loss": 0.1774832308292389, "jvp_rms": 2.067187547683716, "grad_norm": 0.17564058303833008, "grad_norm_postclip": 0.17564058303833008, "clip_fraction": 0.0, "ema_updates": 42150, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3372866999019276, "updates_nfe_5": 11, "loss_nfe_3": 0.40264424681663513, "updates_nfe_3": 9, "loss_nfe_2": 0.3679279110261372, "updates_nfe_2": 14, "loss_nfe_4": 0.33087087869644166, "updates_nfe_4": 10, "loss_nfe_1": 0.40072545905907947, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.4824300911299, "grad_norm": 0.33120907359740953, "grad_rms": 0.0009644484712135507, "grad_to_weight": 0.02009467486081748, "update_to_weight": 7.306840479479682e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.244661833535886, "grad_norm": 0.24401065418607754, "grad_rms": 0.0007105351910607243, "grad_to_weight": 0.015020974686117248, "update_to_weight": 7.878123519641577e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058236858489048, "grad_norm": 0.009202390860888664, "grad_rms": 0.00016474922430517286, "grad_to_weight": 0.0018192882457539802, "update_to_weight": 3.9630966935654306e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.839232901122163, "grad_norm": 0.11039849776982476, "grad_rms": 0.0003532756450565543, "grad_to_weight": 0.008598527546001274, "update_to_weight": 7.696670882732992e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.057599179798749, "grad_norm": 0.039192112823621714, "grad_rms": 0.0008916514407030431, "grad_to_weight": 0.007749153586580034, "update_to_weight": 2.5509622230513197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636010216921004, "grad_norm": 0.21088517387723493, "grad_rms": 0.000508158520764591, "grad_to_weight": 0.012676427288000645, "update_to_weight": 7.300448465915482e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55662501304383, "grad_norm": 0.024401483871660964, "grad_rms": 3.3665514441401645e-05, "grad_to_weight": 0.000993682310117923, "update_to_weight": 8.946302827841843e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824357109450595, "grad_norm": 0.00889337211681509, "grad_rms": 1.0796346566585567e-05, "grad_to_weight": 0.0002885176837666608, "update_to_weight": 8.293334394185e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.597576643029676, "grad_norm": 0.0031416239484296092, "grad_rms": 2.820242848110558e-06, "grad_to_weight": 9.080473990546236e-05, "update_to_weight": 0.00010390434364305669, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.027899044061186, "grad_norm": 0.002363946458129302, "grad_rms": 1.9014682638831111e-06, "grad_to_weight": 6.38423059141523e-05, "update_to_weight": 0.00010902081070631698, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65972775298954, "grad_norm": 0.0010299800170121642, "grad_rms": 1.8608560058658771e-06, "grad_to_weight": 4.1767696193941074e-05, "update_to_weight": 7.026246065511441e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.27822114163141, "grad_norm": 0.002394942907220518, "grad_rms": 1.9264006239360783e-06, "grad_to_weight": 6.42450963022456e-05, "update_to_weight": 0.00011112870498439921, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.38455302753187, "grad_norm": 0.007075977576720941, "grad_rms": 3.1414888749504366e-06, "grad_to_weight": 0.00011163567839071772, "update_to_weight": 0.00011747298216223658, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74874199872471, "grad_norm": 0.02869789020082196, "grad_rms": 2.1413533647775454e-05, "grad_to_weight": 0.0005654896864545553, "update_to_weight": 8.684095554675986e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.05986273739667, "grad_norm": 0.12214183238109699, "grad_rms": 0.00010489058508325576, "grad_to_weight": 0.0030489827981131645, "update_to_weight": 9.437997324535921e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.195799559384174, "grad_norm": 0.29562236635514916, "grad_rms": 0.0004608333768747592, "grad_to_weight": 0.010484624340321628, "update_to_weight": 7.203176302699582e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.405588797619592, "grad_norm": 0.00713228327801499, "grad_rms": 0.0006071402719631676, "grad_to_weight": 0.0008485167963527797, "update_to_weight": 3.2890764269151253e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3068198407581038, "grad_norm": 0.1192600402224126, "grad_rms": 0.0027608266072172453, "grad_to_weight": 0.09125974101619719, "update_to_weight": 7.154738104075162e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42150}
{"step": 42200, "lr": 1.5857872655382737e-05, "window_steps": 50, "elapsed_seconds": 708.685501648, "steps_per_second": 1.6932757862401326, "eta_seconds": 4606.455760712, "images_seen": 43212800, "images_per_second": 1733.9144051098958, "effective_batch_size": 1024, "loss": 0.371521532535553, "edge_loss": 0.1952410787343979, "diag_loss": 0.17628046870231628, "jvp_rms": 2.149648427963257, "grad_norm": 0.29567062854766846, "grad_norm_postclip": 0.28186550736427307, "clip_fraction": 0.05999999865889549, "ema_updates": 42200, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39815575982394974, "updates_nfe_1": 19, "loss_nfe_2": 0.33407230377197267, "updates_nfe_2": 5, "loss_nfe_5": 0.3456992473867204, "updates_nfe_5": 9, "loss_nfe_4": 0.35310671230157215, "updates_nfe_4": 6, "loss_nfe_3": 0.37371120940555225, "updates_nfe_3": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.483619650940597, "grad_norm": 0.13896791665170685, "grad_rms": 0.0004046609994911657, "grad_to_weight": 0.008430667510808343, "update_to_weight": 5.796006574321435e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24594488165143, "grad_norm": 0.13657823562530308, "grad_rms": 0.0003977024817562181, "grad_to_weight": 0.00840691240923499, "update_to_weight": 5.379571904916635e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058151449654951, "grad_norm": 0.005055409830634783, "grad_rms": 9.050635435206846e-05, "grad_to_weight": 0.0009994579800451891, "update_to_weight": 4.0535874147899815e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.840271564594694, "grad_norm": 0.05545670805093494, "grad_rms": 0.0001774616929141041, "grad_to_weight": 0.004318966913741083, "update_to_weight": 8.09687707151401e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.057341785228401, "grad_norm": 0.020380354214948285, "grad_rms": 0.0004636691132110723, "grad_to_weight": 0.004029855026701119, "update_to_weight": 3.959416152734504e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636182652345084, "grad_norm": 0.08828865891568498, "grad_rms": 0.00021274437405922894, "grad_to_weight": 0.005307026302890438, "update_to_weight": 0.00011850338079445365, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55700653223472, "grad_norm": 0.02981913496717732, "grad_rms": 4.113997837785104e-05, "grad_to_weight": 0.0012142821613063985, "update_to_weight": 0.00012333555047640615, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823875942603326, "grad_norm": 0.023394501751922366, "grad_rms": 2.8400380120021304e-05, "grad_to_weight": 0.0007589733943740532, "update_to_weight": 0.00010572998513758922, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.5971708872046, "grad_norm": 0.015826158819378003, "grad_rms": 1.4207178184302887e-05, "grad_to_weight": 0.00045744083731514423, "update_to_weight": 0.00011898155966118929, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.02914518016096, "grad_norm": 0.012948542452466029, "grad_rms": 1.0415313110091016e-05, "grad_to_weight": 0.0003496851571773103, "update_to_weight": 0.00012656019803610824, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660419532195842, "grad_norm": 0.005221057603986585, "grad_rms": 9.43283970453516e-06, "grad_to_weight": 0.00021171811765692558, "update_to_weight": 8.731698518163986e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.279375411273826, "grad_norm": 0.012015599427371163, "grad_rms": 9.664889365031745e-06, "grad_to_weight": 0.00032231225160863264, "update_to_weight": 0.00013038551104172132, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.38613029418877, "grad_norm": 0.024989444443284065, "grad_rms": 1.1094447496277715e-05, "grad_to_weight": 0.0003942415214070119, "update_to_weight": 0.0001428797467412436, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74896088837682, "grad_norm": 0.03143708842699672, "grad_rms": 2.3457443948276916e-05, "grad_to_weight": 0.0006194627018303511, "update_to_weight": 0.00011472826696098942, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.060373139721875, "grad_norm": 0.05995741128648841, "grad_rms": 5.1489058476663806e-05, "grad_to_weight": 0.0014966763059687436, "update_to_weight": 0.00013258892617452062, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.195929440465672, "grad_norm": 0.167194790101874, "grad_rms": 0.00026063298480585726, "grad_to_weight": 0.0059297492020930756, "update_to_weight": 0.00010890627225558122, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.405704523136263, "grad_norm": 0.00489206257476616, "grad_rms": 0.00041644002156501525, "grad_to_weight": 0.0005819931644398174, "update_to_weight": 6.8337223027563e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3068586090680023, "grad_norm": 0.06554575134264476, "grad_rms": 0.0015173603325920324, "grad_to_weight": 0.05015519727064375, "update_to_weight": 0.00020235928558485618, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42200}
{"step": 42250, "lr": 1.5784668743375885e-05, "window_steps": 50, "elapsed_seconds": 736.1441911400001, "steps_per_second": 1.6980368996245665, "eta_seconds": 4564.093985068001, "images_seen": 43264000, "images_per_second": 1738.789785215556, "effective_batch_size": 1024, "loss": 0.3778175413608551, "edge_loss": 0.1965426802635193, "diag_loss": 0.18127483129501343, "jvp_rms": 2.4083592891693115, "grad_norm": 0.25853243470191956, "grad_norm_postclip": 0.25853243470191956, "clip_fraction": 0.0, "ema_updates": 42250, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.36419717832045123, "updates_nfe_2": 11, "loss_nfe_3": 0.40069580995119536, "updates_nfe_3": 13, "loss_nfe_4": 0.3591018617153168, "updates_nfe_4": 8, "loss_nfe_5": 0.32433560490608215, "updates_nfe_5": 6, "loss_nfe_1": 0.40473632514476776, "updates_nfe_1": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48414265979628, "grad_norm": 0.17054129203014604, "grad_rms": 0.0004965995846393517, "grad_to_weight": 0.010345778700768273, "update_to_weight": 8.235657295642742e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24499152834065, "grad_norm": 0.14228011693781892, "grad_rms": 0.0004143058032026026, "grad_to_weight": 0.008758398961895437, "update_to_weight": 7.04966071552713e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058219701421792, "grad_norm": 0.00679193162873953, "grad_rms": 0.00012159508156998123, "grad_to_weight": 0.0013427514085302419, "update_to_weight": 3.244024521594873e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84057162049321, "grad_norm": 0.05912615964500314, "grad_rms": 0.00018920395304522492, "grad_to_weight": 0.004604636101296251, "update_to_weight": 8.288534401999647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.057067027997472, "grad_norm": 0.022564496366283927, "grad_rms": 0.0005133600677330465, "grad_to_weight": 0.004461972966021602, "update_to_weight": 3.5218009684961106e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63676781881118, "grad_norm": 0.11600960192474613, "grad_rms": 0.0002795420210189174, "grad_to_weight": 0.006973085348559963, "update_to_weight": 9.481756121757697e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.557176784814352, "grad_norm": 0.017511804793580648, "grad_rms": 2.4160166663387703e-05, "grad_to_weight": 0.0007131033402996709, "update_to_weight": 0.00010603170146582958, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823969998441857, "grad_norm": 0.007841446826848735, "grad_rms": 9.51933377060015e-06, "grad_to_weight": 0.0002543944478029637, "update_to_weight": 8.837137558288808e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59747608659482, "grad_norm": 0.0034833692990618448, "grad_rms": 3.1270284140524557e-06, "grad_to_weight": 0.00010068275761917545, "update_to_weight": 9.85649430888507e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03005503082513, "grad_norm": 0.0029202935553281325, "grad_rms": 2.348972624816994e-06, "grad_to_weight": 7.886279274759858e-05, "update_to_weight": 0.00010460961226596617, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660243804790714, "grad_norm": 0.0010596980369375504, "grad_rms": 1.914547295936744e-06, "grad_to_weight": 4.2971920526255474e-05, "update_to_weight": 6.571483134061973e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.280303712843796, "grad_norm": 0.003000000810968425, "grad_rms": 2.413086097641235e-06, "grad_to_weight": 8.04714691724699e-05, "update_to_weight": 0.00010730454541168659, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.38739939722839, "grad_norm": 0.008046856361229846, "grad_rms": 3.5725254161760977e-06, "grad_to_weight": 0.00012694725509722826, "update_to_weight": 0.00011763916743619829, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74973541341127, "grad_norm": 0.021578161401876784, "grad_rms": 1.6100998435870512e-05, "grad_to_weight": 0.0004251876630705443, "update_to_weight": 9.819315345301699e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.061198911204905, "grad_norm": 0.06766308325448832, "grad_rms": 5.81063854434094e-05, "grad_to_weight": 0.0016889929680952039, "update_to_weight": 0.00012192374761756881, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.1963201624454, "grad_norm": 0.16208216427313607, "grad_rms": 0.0002526631256426168, "grad_to_weight": 0.005748344583241499, "update_to_weight": 9.49216857314858e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.405695446682731, "grad_norm": 0.004863197820136458, "grad_rms": 0.000413982890476287, "grad_to_weight": 0.0005785598408821362, "update_to_weight": 5.432907958951015e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3068424633469227, "grad_norm": 0.053702557516942716, "grad_rms": 0.0012431946978374968, "grad_to_weight": 0.041093367428088, "update_to_weight": 0.00015035328419219133, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42250}
{"step": 42300, "lr": 1.57118936839488e-05, "window_steps": 50, "elapsed_seconds": 763.593845278, "steps_per_second": 1.702475744191877, "eta_seconds": 4522.825083569693, "images_seen": 43315200, "images_per_second": 1743.335162052482, "effective_batch_size": 1024, "loss": 0.35831549763679504, "edge_loss": 0.18580211699008942, "diag_loss": 0.17251338064670563, "jvp_rms": 1.9698827266693115, "grad_norm": 0.21910293400287628, "grad_norm_postclip": 0.21910293400287628, "clip_fraction": 0.0, "ema_updates": 42300, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4031459490458171, "updates_nfe_1": 9, "loss_nfe_4": 0.34694133400917054, "updates_nfe_4": 10, "loss_nfe_2": 0.36842105431216104, "updates_nfe_2": 14, "loss_nfe_5": 0.3260223902761936, "updates_nfe_5": 8, "loss_nfe_3": 0.339108450545205, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.482904999621958, "grad_norm": 0.026712420258254627, "grad_rms": 7.778395863575476e-05, "grad_to_weight": 0.001620613615067689, "update_to_weight": 8.926849949134808e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24657134050953, "grad_norm": 0.008142593062976843, "grad_rms": 2.37104356653214e-05, "grad_to_weight": 0.0005011883979897922, "update_to_weight": 8.420647473003805e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0583282990197524, "grad_norm": 0.0012625907821514802, "grad_rms": 2.2604000973093905e-05, "grad_to_weight": 0.00024960633385463657, "update_to_weight": 4.241050676605394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.840035674143637, "grad_norm": 0.008473075094236924, "grad_rms": 2.7113875007340377e-05, "grad_to_weight": 0.0006598949807670245, "update_to_weight": 9.16809002900609e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.057240520916656, "grad_norm": 0.01941407639735164, "grad_rms": 0.00044168553166605925, "grad_to_weight": 0.0038388675241083294, "update_to_weight": 2.986005888534071e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636068918542627, "grad_norm": 0.03020677706821032, "grad_rms": 7.278762593800598e-05, "grad_to_weight": 0.0018157400775457072, "update_to_weight": 9.958037395785103e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.556591458955197, "grad_norm": 0.015091603865363679, "grad_rms": 2.0821135736887082e-05, "grad_to_weight": 0.0006145642765848977, "update_to_weight": 0.00013984587063580449, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8244422534192, "grad_norm": 0.010660299654272492, "grad_rms": 1.2941355434072986e-05, "grad_to_weight": 0.0003458391742056581, "update_to_weight": 0.00012449225291709685, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.598661575449704, "grad_norm": 0.007858529133111037, "grad_rms": 7.0546191868073115e-06, "grad_to_weight": 0.00022713390562735647, "update_to_weight": 0.00014900070052380115, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03091376429631, "grad_norm": 0.0067561503311774984, "grad_rms": 5.4343893435401605e-06, "grad_to_weight": 0.00018244622247727256, "update_to_weight": 0.00015011763055021308, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660419532195842, "grad_norm": 0.0025066413721169883, "grad_rms": 4.528727329474676e-06, "grad_to_weight": 0.00010164633934327025, "update_to_weight": 0.00010001654504956122, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.28108300989424, "grad_norm": 0.006877752123724384, "grad_rms": 5.5322011821137705e-06, "grad_to_weight": 0.00018448369973316113, "update_to_weight": 0.00014448458530687088, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.38928663548952, "grad_norm": 0.016824445407597016, "grad_rms": 7.469470826060733e-06, "grad_to_weight": 0.00026541465128553093, "update_to_weight": 0.00015104537174421634, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.74993505614687, "grad_norm": 0.02442687175414413, "grad_rms": 1.8226623509845318e-05, "grad_to_weight": 0.0004813182859666641, "update_to_weight": 0.0001126593367097431, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.061378689456006, "grad_norm": 0.025086605899252608, "grad_rms": 2.1543387054449465e-05, "grad_to_weight": 0.0006262042575647877, "update_to_weight": 0.00011181589010847084, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19584285314448, "grad_norm": 0.06620231687639114, "grad_rms": 0.00010320003056341408, "grad_to_weight": 0.0023479460153470132, "update_to_weight": 8.074641856746797e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406196452257424, "grad_norm": 0.010591225128298683, "grad_rms": 0.000901584955920042, "grad_to_weight": 0.0012599307175904133, "update_to_weight": 4.384302056191806e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3075551024691243, "grad_norm": 0.11509982837366753, "grad_rms": 0.002664519214210704, "grad_to_weight": 0.08802675172642326, "update_to_weight": 0.000107827901273993, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42300}
{"step": 42350, "lr": 1.5639548272957888e-05, "window_steps": 50, "elapsed_seconds": 791.0491039650001, "steps_per_second": 1.7065944367212516, "eta_seconds": 4482.611589135, "images_seen": 43366400, "images_per_second": 1747.5527032025616, "effective_batch_size": 1024, "loss": 0.36100679636001587, "edge_loss": 0.1889551728963852, "diag_loss": 0.17205165326595306, "jvp_rms": 1.8283593654632568, "grad_norm": 0.2706964910030365, "grad_norm_postclip": 0.2706964910030365, "clip_fraction": 0.0, "ema_updates": 42350, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.35199973529035394, "updates_nfe_2": 11, "loss_nfe_4": 0.3466060360272725, "updates_nfe_4": 9, "loss_nfe_1": 0.40984322875738144, "updates_nfe_1": 12, "loss_nfe_5": 0.3942989110946655, "updates_nfe_5": 6, "loss_nfe_3": 0.31458163013060886, "updates_nfe_3": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48334471750718, "grad_norm": 0.03272301835085983, "grad_rms": 9.528623319160952e-05, "grad_to_weight": 0.001985217133516857, "update_to_weight": 6.622543620923666e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.247113250200595, "grad_norm": 0.008511082024992522, "grad_rms": 2.478343953026705e-05, "grad_to_weight": 0.000523851954124062, "update_to_weight": 7.650776556248318e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058048127548383, "grad_norm": 0.0010419251437934604, "grad_rms": 1.8653452327654298e-05, "grad_to_weight": 0.00020599352112105697, "update_to_weight": 2.917072142901142e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.841563834625049, "grad_norm": 0.010488978894169856, "grad_rms": 3.356477542428357e-05, "grad_to_weight": 0.0008167991865514186, "update_to_weight": 9.653238711303576e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.056440895352999, "grad_norm": 0.028459098921782156, "grad_rms": 0.0006474669194007627, "grad_to_weight": 0.005628286676491524, "update_to_weight": 5.331521087020247e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636520180342327, "grad_norm": 0.06576723242186727, "grad_rms": 0.000158475718931917, "grad_to_weight": 0.003953184422519902, "update_to_weight": 9.578325377430139e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.557335851851377, "grad_norm": 0.019008870346342376, "grad_rms": 2.622559359608162e-05, "grad_to_weight": 0.0007740607719427878, "update_to_weight": 0.00010891179684928972, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824173950116347, "grad_norm": 0.009431858082010445, "grad_rms": 1.1450055983568042e-05, "grad_to_weight": 0.00030598899737830103, "update_to_weight": 9.292322676515859e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59633289815555, "grad_norm": 0.006156824482246186, "grad_rms": 5.526995114041684e-06, "grad_to_weight": 0.00017796176549608896, "update_to_weight": 0.00010627517987205537, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03035995765839, "grad_norm": 0.005182043870374947, "grad_rms": 4.1682382135530615e-06, "grad_to_weight": 0.00013994041311778363, "update_to_weight": 0.00011231854033849567, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.65981314360155, "grad_norm": 0.0017585804172743312, "grad_rms": 3.177212059682524e-06, "grad_to_weight": 7.131361486940658e-05, "update_to_weight": 7.251349965827747e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.28169694042658, "grad_norm": 0.005172458538903477, "grad_rms": 4.1605281389322415e-06, "grad_to_weight": 0.00013873989017100502, "update_to_weight": 0.0001158715332103586, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.389558162085144, "grad_norm": 0.012736018536814676, "grad_rms": 5.654350951618871e-06, "grad_to_weight": 0.00020091666366011055, "update_to_weight": 0.0001278148297822932, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75019242574357, "grad_norm": 0.02450888933706239, "grad_rms": 1.8287822652334166e-05, "grad_to_weight": 0.00048293194893641423, "update_to_weight": 0.00010026329441974561, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.06230956224238, "grad_norm": 0.027182414941667468, "grad_rms": 2.3343185144883996e-05, "grad_to_weight": 0.0006785034422300541, "update_to_weight": 0.00012741349009122343, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19622275304021, "grad_norm": 0.07699273791472058, "grad_rms": 0.00012002076786520516, "grad_to_weight": 0.002730604683792937, "update_to_weight": 0.0001107605971782412, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.405702707846341, "grad_norm": 0.011120589085639083, "grad_rms": 0.000946647408503473, "grad_to_weight": 0.0013229814891333855, "update_to_weight": 9.805295430833287e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3071125814667466, "grad_norm": 0.11862280364108584, "grad_rms": 0.0027460748118502635, "grad_to_weight": 0.09075178781308643, "update_to_weight": 0.00037511382160604123, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42350}
{"step": 42400, "lr": 1.5567633301560955e-05, "window_steps": 50, "elapsed_seconds": 818.514535117, "steps_per_second": 1.7104155637259164, "eta_seconds": 4443.364619206572, "images_seen": 43417600, "images_per_second": 1751.4655372553384, "effective_batch_size": 1024, "loss": 0.36354154348373413, "edge_loss": 0.1876566857099533, "diag_loss": 0.17588481307029724, "jvp_rms": 2.0893750190734863, "grad_norm": 0.23112161457538605, "grad_norm_postclip": 0.23112161457538605, "clip_fraction": 0.0, "ema_updates": 42400, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.37989577816592324, "updates_nfe_5": 9, "loss_nfe_4": 0.34174502747399466, "updates_nfe_4": 14, "loss_nfe_3": 0.3585787063295191, "updates_nfe_3": 11, "loss_nfe_1": 0.4058155980375078, "updates_nfe_1": 9, "loss_nfe_2": 0.3395538457802364, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.483712220088275, "grad_norm": 0.0067422820930300225, "grad_rms": 1.9632866897291918e-05, "grad_to_weight": 0.0004090269232444726, "update_to_weight": 3.4578088219535884e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.244966167439376, "grad_norm": 0.00704460620182798, "grad_rms": 2.0513205172370713e-05, "grad_to_weight": 0.00043364856098917874, "update_to_weight": 3.277658947841139e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058883317001793, "grad_norm": 0.0008938430816893375, "grad_rms": 1.6002358146374462e-05, "grad_to_weight": 0.00017668782331573606, "update_to_weight": 2.616887742818501e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.840290578245847, "grad_norm": 0.004325427053201442, "grad_rms": 1.384138428722784e-05, "grad_to_weight": 0.0003368636423641086, "update_to_weight": 5.718146986395793e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.056948973977792, "grad_norm": 0.010837283696691518, "grad_rms": 0.00024655674127470106, "grad_to_weight": 0.002143047863931069, "update_to_weight": 2.2027852467683633e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63655228183327, "grad_norm": 0.040119664253642245, "grad_rms": 9.667417043063964e-05, "grad_to_weight": 0.0024115371727261056, "update_to_weight": 7.276682639676309e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55714447419656, "grad_norm": 0.013419184346850011, "grad_rms": 1.851378165347447e-05, "grad_to_weight": 0.0005464472614456551, "update_to_weight": 8.892084414543164e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823762085151632, "grad_norm": 0.009360798790619665, "grad_rms": 1.1363791659242705e-05, "grad_to_weight": 0.00030368774469385527, "update_to_weight": 8.214967106039318e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59688861947614, "grad_norm": 0.006414263153810816, "grad_rms": 5.758098385542449e-06, "grad_to_weight": 0.00018540000010867856, "update_to_weight": 9.583019279993054e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.031043973582996, "grad_norm": 0.005297678252980358, "grad_rms": 4.261250095434643e-06, "grad_to_weight": 0.0001430604618319588, "update_to_weight": 0.00010099590630018066, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660071788723567, "grad_norm": 0.0018461297274809218, "grad_rms": 3.3353866427001133e-06, "grad_to_weight": 7.486311245554082e-05, "update_to_weight": 6.451381966370549e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.282857650794085, "grad_norm": 0.005368304646048031, "grad_rms": 4.31805926915721e-06, "grad_to_weight": 0.0001439885508865679, "update_to_weight": 9.952262365083479e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.39171877907733, "grad_norm": 0.013285327712265989, "grad_rms": 5.898225193005059e-06, "grad_to_weight": 0.00020957513012962917, "update_to_weight": 0.00010761483932822649, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75069513455579, "grad_norm": 0.019441670388015326, "grad_rms": 1.4506810783281965e-05, "grad_to_weight": 0.000383081854080411, "update_to_weight": 7.88024373380476e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.062917436397775, "grad_norm": 0.020932773682267584, "grad_rms": 1.7976239885592464e-05, "grad_to_weight": 0.0005224974869965371, "update_to_weight": 8.991617358285262e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.196171883550413, "grad_norm": 0.03427845275204298, "grad_rms": 5.343525028410186e-05, "grad_to_weight": 0.0012157130015241875, "update_to_weight": 6.312286777391917e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406245008275702, "grad_norm": 0.004326817956242297, "grad_rms": 0.0003683232042655477, "grad_to_weight": 0.0005147147093598475, "update_to_weight": 4.698571129223788e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3077408017794168, "grad_norm": 0.05395300459869035, "grad_rms": 0.001248992456799336, "grad_to_weight": 0.04125665003743676, "update_to_weight": 0.00017750579074864132, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42400}
{"step": 42450, "lr": 1.549614955620861e-05, "window_steps": 50, "elapsed_seconds": 846.0041050970001, "steps_per_second": 1.7139396738905277, "eta_seconds": 4405.055857574035, "images_seen": 43468800, "images_per_second": 1755.0742260639004, "effective_batch_size": 1024, "loss": 0.36973920464515686, "edge_loss": 0.1935407817363739, "diag_loss": 0.17619851231575012, "jvp_rms": 2.1830859184265137, "grad_norm": 0.19060178101062775, "grad_norm_postclip": 0.19060178101062775, "clip_fraction": 0.0, "ema_updates": 42450, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.37015883922576903, "updates_nfe_2": 10, "loss_nfe_5": 0.3341435581445694, "updates_nfe_5": 10, "loss_nfe_1": 0.4027914941310883, "updates_nfe_1": 15, "loss_nfe_4": 0.38240694254636765, "updates_nfe_4": 4, "loss_nfe_3": 0.352039944041859, "updates_nfe_3": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48375202466195, "grad_norm": 0.04577115347099455, "grad_rms": 0.00013328112817476495, "grad_to_weight": 0.002776743632306203, "update_to_weight": 4.304013712164861e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245430171738874, "grad_norm": 0.03772103822766973, "grad_rms": 0.00010983997888742777, "grad_to_weight": 0.002321947638745238, "update_to_weight": 5.005231031945429e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.058491756633464, "grad_norm": 0.00425565471007698, "grad_rms": 7.61884409165524e-05, "grad_to_weight": 0.0008412892448616364, "update_to_weight": 3.6580767347140845e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.840667874563369, "grad_norm": 0.026417621089568865, "grad_rms": 8.453649569340411e-05, "grad_to_weight": 0.0020573401124952904, "update_to_weight": 7.45757030473577e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.056634401249067, "grad_norm": 0.04606221785387768, "grad_rms": 0.0010479517421330001, "grad_to_weight": 0.009109264027966824, "update_to_weight": 2.9817525613871494e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.6364550599847, "grad_norm": 0.12355809507636893, "grad_rms": 0.0002977312139498689, "grad_to_weight": 0.007426948507411323, "update_to_weight": 9.01374919105082e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.558167208897252, "grad_norm": 0.0428892641294574, "grad_rms": 5.917218594268443e-05, "grad_to_weight": 0.001746435870585608, "update_to_weight": 0.00011619705498276622, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823940296629114, "grad_norm": 0.029934937426232735, "grad_rms": 3.634031666026837e-05, "grad_to_weight": 0.0009711586882844566, "update_to_weight": 9.601179283718009e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.596613406223554, "grad_norm": 0.02049106385133188, "grad_rms": 1.8394873869542108e-05, "grad_to_weight": 0.0005922852508923825, "update_to_weight": 0.00011497241904485608, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03083794601973, "grad_norm": 0.016437948809879015, "grad_rms": 1.3222058341395171e-05, "grad_to_weight": 0.0004438989156508097, "update_to_weight": 0.00011940301498990482, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660196778934697, "grad_norm": 0.006115002551675407, "grad_rms": 1.1047922324920168e-05, "grad_to_weight": 0.0002479705497280939, "update_to_weight": 7.608604978553683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.28385462139997, "grad_norm": 0.015607214188957219, "grad_rms": 1.2553847133836017e-05, "grad_to_weight": 0.00041860516696680503, "update_to_weight": 0.00012040129905331381, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.393723349841586, "grad_norm": 0.028132225141720753, "grad_rms": 1.2489733235032565e-05, "grad_to_weight": 0.00044376988217700343, "update_to_weight": 0.000133259639749251, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75098136366429, "grad_norm": 0.03642483095100467, "grad_rms": 2.717915281317548e-05, "grad_to_weight": 0.0007177167804893605, "update_to_weight": 0.00010027465678943665, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.063433893658726, "grad_norm": 0.0431087027232376, "grad_rms": 3.7020052529687795e-05, "grad_to_weight": 0.001076011178613945, "update_to_weight": 0.00011537555851780792, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19633098569187, "grad_norm": 0.09344529762870722, "grad_rms": 0.00014566797698781092, "grad_to_weight": 0.003314094222972688, "update_to_weight": 8.650704373666372e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406331682340141, "grad_norm": 0.007408340237341415, "grad_rms": 0.0006306398009119698, "grad_to_weight": 0.00088128098168012, "update_to_weight": 5.5206235086446804e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.307841161416119, "grad_norm": 0.10940056740430161, "grad_rms": 0.00253258339315651, "grad_to_weight": 0.08364973563444328, "update_to_weight": 0.00020377599782309027, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42450}
{"step": 42500, "lr": 1.5425097818635637e-05, "window_steps": 50, "elapsed_seconds": 873.781596037, "steps_per_second": 1.7166761199860325, "eta_seconds": 4368.907980185, "images_seen": 43520000, "images_per_second": 1757.8763468656973, "effective_batch_size": 1024, "loss": 0.3630458414554596, "edge_loss": 0.18846014142036438, "diag_loss": 0.1745857447385788, "jvp_rms": 1.9151562452316284, "grad_norm": 0.34880736470222473, "grad_norm_postclip": 0.33721357583999634, "clip_fraction": 0.03999999910593033, "ema_updates": 42500, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.36960374457495554, "updates_nfe_4": 7, "loss_nfe_1": 0.4022623121738434, "updates_nfe_1": 10, "loss_nfe_3": 0.3418373540043831, "updates_nfe_3": 8, "loss_nfe_5": 0.3483955721060435, "updates_nfe_5": 15, "loss_nfe_2": 0.35818096101284025, "updates_nfe_2": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.483129951500317, "grad_norm": 0.06583034908622024, "grad_rms": 0.00019169154650887572, "grad_to_weight": 0.003993801497647494, "update_to_weight": 9.376286648948505e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245596421066598, "grad_norm": 0.01988795263899886, "grad_rms": 5.791177551362823e-05, "grad_to_weight": 0.001224205755426068, "update_to_weight": 0.0001009880838211575, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.059631099899026, "grad_norm": 0.0031181889401437142, "grad_rms": 5.582453700255422e-05, "grad_to_weight": 0.0006162878041061023, "update_to_weight": 4.127328767205723e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.840603705263433, "grad_norm": 0.013451104526031293, "grad_rms": 4.304358957913006e-05, "grad_to_weight": 0.001047544557466376, "update_to_weight": 0.0001060583975567693, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.056763023878328, "grad_norm": 0.06278145641555956, "grad_rms": 0.0014283275901530229, "grad_to_weight": 0.012415344780663418, "update_to_weight": 4.118005203826446e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636267952186667, "grad_norm": 0.11220107066870579, "grad_rms": 0.0002703648106262996, "grad_to_weight": 0.006744365442488446, "update_to_weight": 0.0001136784758017393, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55797086702134, "grad_norm": 0.034954916779260374, "grad_rms": 4.822556127403583e-05, "grad_to_weight": 0.001423363394660631, "update_to_weight": 0.00014154038641451042, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.823972968621558, "grad_norm": 0.017031262675551852, "grad_rms": 2.06755561216369e-05, "grad_to_weight": 0.0005525330136024152, "update_to_weight": 0.00010976120039709725, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.598598068147936, "grad_norm": 0.011139030325316117, "grad_rms": 9.999532447402742e-06, "grad_to_weight": 0.000321950337507198, "update_to_weight": 0.00012527107970254774, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03187631326895, "grad_norm": 0.00904964186837453, "grad_rms": 7.2791863593387775e-06, "grad_to_weight": 0.00024437438146043757, "update_to_weight": 0.00012996792400992826, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660396019446594, "grad_norm": 0.0035699725182708275, "grad_rms": 6.4498385324711994e-06, "grad_to_weight": 0.00014476541720804617, "update_to_weight": 8.968989703415139e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.2851658672722, "grad_norm": 0.009150978113920752, "grad_rms": 7.360697365742742e-06, "grad_to_weight": 0.00024543214173959745, "update_to_weight": 0.00012642056906887777, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.396016685913516, "grad_norm": 0.02025997820476683, "grad_rms": 8.994728353351787e-06, "grad_to_weight": 0.0003195780943957724, "update_to_weight": 0.0001408106737627373, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.751176190957636, "grad_norm": 0.03057384168206671, "grad_rms": 2.2813314254786e-05, "grad_to_weight": 0.0006024262682509823, "update_to_weight": 0.00011340276289762335, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.063788858948115, "grad_norm": 0.031745807048882524, "grad_rms": 2.7262046183390707e-05, "grad_to_weight": 0.0007923815483515409, "update_to_weight": 0.00012875085995122827, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.196029015556316, "grad_norm": 0.0684139510686685, "grad_rms": 0.00010664765486128077, "grad_to_weight": 0.002426368302817505, "update_to_weight": 0.00010578703060814766, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40651682603945, "grad_norm": 0.007011357017679174, "grad_rms": 0.000596846345078058, "grad_to_weight": 0.0008340383018042949, "update_to_weight": 1.0589304661101024e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3082564189599601, "grad_norm": 0.06616558614794645, "grad_rms": 0.0015317092831656262, "grad_to_weight": 0.050575395762664684, "update_to_weight": 0.00023211237671770692, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42500}
{"step": 42550, "lr": 1.5354478865852468e-05, "window_steps": 50, "elapsed_seconds": 903.5652241350001, "steps_per_second": 1.7154267988609722, "eta_seconds": 4342.942528906936, "images_seen": 43571200, "images_per_second": 1756.5970420336355, "effective_batch_size": 1024, "loss": 0.34889596700668335, "edge_loss": 0.18189097940921783, "diag_loss": 0.16700495779514313, "jvp_rms": 1.5039843320846558, "grad_norm": 0.2377571314573288, "grad_norm_postclip": 0.2377571314573288, "clip_fraction": 0.0, "ema_updates": 42550, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3453151099383831, "updates_nfe_2": 8, "loss_nfe_5": 0.32273579239845274, "updates_nfe_5": 10, "loss_nfe_1": 0.4169178174601661, "updates_nfe_1": 9, "loss_nfe_3": 0.3143217048861764, "updates_nfe_3": 11, "loss_nfe_4": 0.35376014312108356, "updates_nfe_4": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.483527081273063, "grad_norm": 0.06990537062383034, "grad_rms": 0.0002035576111955199, "grad_to_weight": 0.004240923091226625, "update_to_weight": 6.0003495213463065e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25074179496797, "grad_norm": 0.05541088854064147, "grad_rms": 0.00016135109512900765, "grad_to_weight": 0.003409745182081437, "update_to_weight": 7.000450816697907e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.05993813578667, "grad_norm": 0.004385430047638292, "grad_rms": 7.851179215433099e-05, "grad_to_weight": 0.0008666963765074744, "update_to_weight": 3.563612573432963e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.843190424625927, "grad_norm": 0.025913389866169517, "grad_rms": 8.292295371319113e-05, "grad_to_weight": 0.002017675438065793, "update_to_weight": 9.30417912725336e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.055679247580779, "grad_norm": 0.024426296268431343, "grad_rms": 0.0005557174821577666, "grad_to_weight": 0.004831456876960639, "update_to_weight": 4.4629010856353945e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636472486584257, "grad_norm": 0.12020260789641725, "grad_rms": 0.0002896456791990883, "grad_to_weight": 0.00722524609669203, "update_to_weight": 0.00010156627140818047, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.558791018871304, "grad_norm": 0.01970630738348573, "grad_rms": 2.718781281067563e-05, "grad_to_weight": 0.0008024135784348317, "update_to_weight": 0.0001106263057970867, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824109596578374, "grad_norm": 0.007700987600024851, "grad_rms": 9.348819541425135e-06, "grad_to_weight": 0.00024983649814428697, "update_to_weight": 8.510359850414772e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.59780068877211, "grad_norm": 0.0025262507727820734, "grad_rms": 2.267820971390856e-06, "grad_to_weight": 7.301766940353257e-05, "update_to_weight": 9.193116380076456e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03208727956406, "grad_norm": 0.0020477975922487367, "grad_rms": 1.647170188278536e-06, "grad_to_weight": 5.529792519631487e-05, "update_to_weight": 9.171686261969701e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660419532195842, "grad_norm": 0.000803266691824077, "grad_rms": 1.4512549982561409e-06, "grad_to_weight": 3.257311542390259e-05, "update_to_weight": 5.6066144522819034e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.28572734797039, "grad_norm": 0.0021339797399119935, "grad_rms": 1.7164918170030083e-06, "grad_to_weight": 5.723315305067139e-05, "update_to_weight": 9.551939133724265e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.39618805698218, "grad_norm": 0.006466676447401428, "grad_rms": 2.8709802846534336e-06, "grad_to_weight": 0.00010200418425140967, "update_to_weight": 0.00010878080268210435, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.751621164069725, "grad_norm": 0.021749303683582678, "grad_rms": 1.622869984465774e-05, "grad_to_weight": 0.0004285440185895063, "update_to_weight": 9.677603906879672e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.06437842953903, "grad_norm": 0.04485995752414112, "grad_rms": 3.852396103601752e-05, "grad_to_weight": 0.0011196968300166204, "update_to_weight": 0.00012866073702487579, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.196757418296347, "grad_norm": 0.13478257584046868, "grad_rms": 0.0002101069358663553, "grad_to_weight": 0.004780073603534668, "update_to_weight": 0.00010750325758515166, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40587742770437, "grad_norm": 0.0060111229208096905, "grad_rms": 0.0005117007643532853, "grad_to_weight": 0.0007151095138501582, "update_to_weight": 4.380076802506961e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3075864188553517, "grad_norm": 0.0923200975099068, "grad_rms": 0.0021371767199718046, "grad_to_weight": 0.07060343865510847, "update_to_weight": 0.00010833930694777782, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42550}
{"step": 42600, "lr": 1.5284293470136648e-05, "window_steps": 50, "elapsed_seconds": 931.021677012, "steps_per_second": 1.7185421558980278, "eta_seconds": 4305.975256180501, "images_seen": 43622400, "images_per_second": 1759.7871676395805, "effective_batch_size": 1024, "loss": 0.3651902675628662, "edge_loss": 0.18792860209941864, "diag_loss": 0.17726176977157593, "jvp_rms": 2.1302733421325684, "grad_norm": 0.20271675288677216, "grad_norm_postclip": 0.20271675288677216, "clip_fraction": 0.0, "ema_updates": 42600, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.395471416413784, "updates_nfe_1": 8, "loss_nfe_2": 0.34075360000133514, "updates_nfe_2": 14, "loss_nfe_5": 0.36117842678840345, "updates_nfe_5": 13, "loss_nfe_4": 0.4028165141741435, "updates_nfe_4": 6, "loss_nfe_3": 0.3569975793361664, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48614012320984, "grad_norm": 0.15680747718110366, "grad_rms": 0.00045660805725991437, "grad_to_weight": 0.009511473031843512, "update_to_weight": 6.570524607786e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245892284306247, "grad_norm": 0.05252788657050277, "grad_rms": 0.00015295607499140817, "grad_to_weight": 0.0032333026497562973, "update_to_weight": 9.068530979676322e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060093060267217, "grad_norm": 0.004168871621965973, "grad_rms": 7.463477441126823e-05, "grad_to_weight": 0.00082387252019152, "update_to_weight": 3.665747197714235e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.841727809937852, "grad_norm": 0.025852433871068656, "grad_rms": 8.272789427919214e-05, "grad_to_weight": 0.002013158529264433, "update_to_weight": 8.114096815188346e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.056508604398634, "grad_norm": 0.029035933028125176, "grad_rms": 0.0006605903497267085, "grad_to_weight": 0.005742288859721696, "update_to_weight": 2.0987854444621204e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.6365981409986, "grad_norm": 0.05541813003063228, "grad_rms": 0.000133538050409839, "grad_to_weight": 0.0033310974732305362, "update_to_weight": 6.597078368788403e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.558162238236218, "grad_norm": 0.025123956459351435, "grad_rms": 3.46622739607134e-05, "grad_to_weight": 0.0010230389479321176, "update_to_weight": 8.219375447576315e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824391761094105, "grad_norm": 0.013823900779765596, "grad_rms": 1.67818935000204e-05, "grad_to_weight": 0.00044847278372622525, "update_to_weight": 7.301692495819704e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.599654743919594, "grad_norm": 0.00957639917891258, "grad_rms": 8.59675497077882e-06, "grad_to_weight": 0.00027677730456531505, "update_to_weight": 8.845947588007662e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.033257461410614, "grad_norm": 0.007809154438635294, "grad_rms": 6.2813856387329185e-06, "grad_to_weight": 0.00021086868868536849, "update_to_weight": 9.584703064819414e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660633620933943, "grad_norm": 0.0032649485578211354, "grad_rms": 5.8987543761179975e-06, "grad_to_weight": 0.00013239516096819114, "update_to_weight": 6.117175224684126e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.28663584839245, "grad_norm": 0.008055545541003788, "grad_rms": 6.4795732330609275e-06, "grad_to_weight": 0.00021604377433667265, "update_to_weight": 9.937477217214627e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.39865459374414, "grad_norm": 0.02034761738498826, "grad_rms": 9.033637112839827e-06, "grad_to_weight": 0.0003209471480960427, "update_to_weight": 0.00011611537829034577, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75222247288782, "grad_norm": 0.030902218126208737, "grad_rms": 2.305833923699082e-05, "grad_to_weight": 0.0006088840373979072, "update_to_weight": 8.799179075311058e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0649679914542, "grad_norm": 0.033467711165411605, "grad_rms": 2.8740749480361663e-05, "grad_to_weight": 0.00083533602653945, "update_to_weight": 9.564178665129654e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.196799628279514, "grad_norm": 0.059791969433266035, "grad_rms": 9.320720730183828e-05, "grad_to_weight": 0.0021205232587211305, "update_to_weight": 6.837425152218643e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406643429309858, "grad_norm": 0.006064835710923589, "grad_rms": 0.0005162731006902584, "grad_to_weight": 0.0007214336806267375, "update_to_weight": 2.907211666733174e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3082962381162022, "grad_norm": 0.07079177996538151, "grad_rms": 0.0016388039894687555, "grad_to_weight": 0.05410990103228732, "update_to_weight": 7.529511882165835e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42600}
{"step": 42650, "lr": 1.5214542399024417e-05, "window_steps": 50, "elapsed_seconds": 958.4747925409999, "steps_per_second": 1.7214850227054033, "eta_seconds": 4269.569530409909, "images_seen": 43673600, "images_per_second": 1762.800663250333, "effective_batch_size": 1024, "loss": 0.377536803483963, "edge_loss": 0.19659219682216644, "diag_loss": 0.18094463646411896, "jvp_rms": 2.2756640911102295, "grad_norm": 0.1986880600452423, "grad_norm_postclip": 0.1986880600452423, "clip_fraction": 0.0, "ema_updates": 42650, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39454676849501474, "updates_nfe_1": 14, "loss_nfe_5": 0.3975186305386679, "updates_nfe_5": 7, "loss_nfe_2": 0.36021966735521954, "updates_nfe_2": 12, "loss_nfe_4": 0.36280616905008045, "updates_nfe_4": 14, "loss_nfe_3": 0.38954392075538635, "updates_nfe_3": 3, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.4857088100786, "grad_norm": 0.39704978728452606, "grad_rms": 0.001156170198427866, "grad_to_weight": 0.024084483831340524, "update_to_weight": 5.508418578158693e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.247320805277692, "grad_norm": 0.2466428700545185, "grad_rms": 0.0007181999465659062, "grad_to_weight": 0.015180525639304195, "update_to_weight": 7.388922065813071e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060119069021576, "grad_norm": 0.010621422243501242, "grad_rms": 0.00019015396130061307, "grad_to_weight": 0.0020990459114937385, "update_to_weight": 3.4366114300195274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84273122195033, "grad_norm": 0.1277339763226835, "grad_rms": 0.0004087492474319587, "grad_to_weight": 0.009946013360800171, "update_to_weight": 7.686693980927163e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.055631145615508, "grad_norm": 0.04586519904578569, "grad_rms": 0.001043469409045437, "grad_to_weight": 0.009072101528918353, "update_to_weight": 2.8886089356317163e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636415620770936, "grad_norm": 0.24292739779983052, "grad_rms": 0.0005853689230472699, "grad_to_weight": 0.014602147682372772, "update_to_weight": 7.567383391771086e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.558568586455216, "grad_norm": 0.035404399118714076, "grad_rms": 4.884568971661827e-05, "grad_to_weight": 0.0014416312169855314, "update_to_weight": 8.396357720411622e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.825000633617186, "grad_norm": 0.012265110134143175, "grad_rms": 1.4889557970388e-05, "grad_to_weight": 0.00039789488668386493, "update_to_weight": 7.205552083249569e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.599986381748586, "grad_norm": 0.0041264234673255235, "grad_rms": 3.7042995797818437e-06, "grad_to_weight": 0.00011926084079334211, "update_to_weight": 8.747642998282463e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.033282183163465, "grad_norm": 0.0031359221720824284, "grad_rms": 2.5224160503790805e-06, "grad_to_weight": 8.467848344017752e-05, "update_to_weight": 9.250771097707368e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660435619853462, "grad_norm": 0.0013113562832857129, "grad_rms": 2.3692160772798308e-06, "grad_to_weight": 5.317652548805645e-05, "update_to_weight": 5.90593341356857e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.28738227541758, "grad_norm": 0.003437532401617984, "grad_rms": 2.7650198020639635e-06, "grad_to_weight": 9.219023143612426e-05, "update_to_weight": 9.858735243357963e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.3994228396935, "grad_norm": 0.008682876959896188, "grad_rms": 3.8548965250844925e-06, "grad_to_weight": 0.00013695514203419465, "update_to_weight": 0.00010620281308191469, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75280212784438, "grad_norm": 0.03892252250224297, "grad_rms": 2.9042857834692147e-05, "grad_to_weight": 0.0007669039121071308, "update_to_weight": 8.14705317558681e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.06545700134422, "grad_norm": 0.12840122940560791, "grad_rms": 0.00011026590820859435, "grad_to_weight": 0.0032047863425419053, "update_to_weight": 9.406902221306274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.196674080451736, "grad_norm": 0.30861637177424844, "grad_rms": 0.00048108918995900537, "grad_to_weight": 0.010945133844285799, "update_to_weight": 7.557304770865554e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406311715619179, "grad_norm": 0.006609485655096723, "grad_rms": 0.0005626367828857352, "grad_to_weight": 0.0007862527442107698, "update_to_weight": 4.883679494638804e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.307995104147705, "grad_norm": 0.11575844770546745, "grad_rms": 0.0026797660124833673, "grad_to_weight": 0.08850067354104979, "update_to_weight": 0.00011507077324323135, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42650}
{"step": 42700, "lr": 1.5145226415302315e-05, "window_steps": 50, "elapsed_seconds": 985.9669825370001, "steps_per_second": 1.7241956679175152, "eta_seconds": 4233.858219129471, "images_seen": 43724800, "images_per_second": 1765.5763639475356, "effective_batch_size": 1024, "loss": 0.36722517013549805, "edge_loss": 0.18961800634860992, "diag_loss": 0.1776072233915329, "jvp_rms": 1.9914453029632568, "grad_norm": 0.25828173756599426, "grad_norm_postclip": 0.25588497519493103, "clip_fraction": 0.019999999552965164, "ema_updates": 42700, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39750003814697266, "updates_nfe_1": 9, "loss_nfe_4": 0.35551848778357875, "updates_nfe_4": 13, "loss_nfe_5": 0.32442737902913776, "updates_nfe_5": 7, "loss_nfe_3": 0.37405935471708124, "updates_nfe_3": 11, "loss_nfe_2": 0.3776376962661743, "updates_nfe_2": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.488463096456357, "grad_norm": 0.11737781012566416, "grad_rms": 0.0003417927181176616, "grad_to_weight": 0.007118784172849353, "update_to_weight": 3.9143687113613356e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.245605813633208, "grad_norm": 0.09111082619551673, "grad_rms": 0.00026530582656101757, "grad_to_weight": 0.0056083366321160585, "update_to_weight": 5.1253060076873196e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.059910241337352, "grad_norm": 0.005055965347483797, "grad_rms": 9.051629969902742e-05, "grad_to_weight": 0.000999220362878905, "update_to_weight": 3.7247700007029614e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.842889241800243, "grad_norm": 0.039176897087926967, "grad_rms": 0.00012536623115024485, "grad_to_weight": 0.00305047379529027, "update_to_weight": 6.840434385752945e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.056113465001718, "grad_norm": 0.04364886696728646, "grad_rms": 0.0009930461083225651, "grad_to_weight": 0.008632889129055894, "update_to_weight": 3.1062362428021597e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.637132849536457, "grad_norm": 0.11468669490245781, "grad_rms": 0.00027635428399978265, "grad_to_weight": 0.006893417029224072, "update_to_weight": 0.00010801000993741818, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.559096704826732, "grad_norm": 0.04673347865328861, "grad_rms": 6.447585764759196e-05, "grad_to_weight": 0.0019028989223412206, "update_to_weight": 0.00014242429052366366, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824830348529495, "grad_norm": 0.03529120576220286, "grad_rms": 4.2842701638562925e-05, "grad_to_weight": 0.001144895376979307, "update_to_weight": 0.00013364247348033332, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.6002774442453, "grad_norm": 0.023261022999749786, "grad_rms": 2.088147239505594e-05, "grad_to_weight": 0.0006722785109810888, "update_to_weight": 0.00015819331143531536, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03364476697591, "grad_norm": 0.01993153293567444, "grad_rms": 1.603216401005853e-05, "grad_to_weight": 0.0005382006837590025, "update_to_weight": 0.00016134950119301922, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66016831587302, "grad_norm": 0.007624615603040648, "grad_rms": 1.3775327193720983e-05, "grad_to_weight": 0.00030918749237136833, "update_to_weight": 0.00011268825731332649, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.288053392222686, "grad_norm": 0.017202648753270565, "grad_rms": 1.3837153775875908e-05, "grad_to_weight": 0.0004613447790454674, "update_to_weight": 0.00014372729661494014, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.40031622678431, "grad_norm": 0.029614673951221453, "grad_rms": 1.31478891424299e-05, "grad_to_weight": 0.0004671060921098425, "update_to_weight": 0.00014495863294026453, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75351165498305, "grad_norm": 0.0335945589500436, "grad_rms": 2.506728590237377e-05, "grad_to_weight": 0.0006619159513220646, "update_to_weight": 9.848254352273658e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.066090724908044, "grad_norm": 0.04609634717397039, "grad_rms": 3.958572367076201e-05, "grad_to_weight": 0.001150507732098592, "update_to_weight": 0.00010375003679502386, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19685374355012, "grad_norm": 0.12299857629750652, "grad_rms": 0.00019173735047459936, "grad_to_weight": 0.004362138322813474, "update_to_weight": 7.088478820656196e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40670468827235, "grad_norm": 0.004137965934368666, "grad_rms": 0.00035224705256885544, "grad_to_weight": 0.0004922221117319935, "update_to_weight": 2.0419925961117346e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.308468986371088, "grad_norm": 0.04940597501921934, "grad_rms": 0.0011437303738468583, "grad_to_weight": 0.037758613718649935, "update_to_weight": 8.299660924026977e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42700}
{"step": 42750, "lr": 1.5076346276998838e-05, "window_steps": 50, "elapsed_seconds": 1013.4427844280001, "steps_per_second": 1.7267871722898713, "eta_seconds": 4198.548678344572, "images_seen": 43776000, "images_per_second": 1768.2300644248282, "effective_batch_size": 1024, "loss": 0.3368394076824188, "edge_loss": 0.17600253224372864, "diag_loss": 0.160836860537529, "jvp_rms": 1.5903124809265137, "grad_norm": 0.19668997824192047, "grad_norm_postclip": 0.19668997824192047, "clip_fraction": 0.0, "ema_updates": 42750, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.38124795470918926, "updates_nfe_2": 7, "loss_nfe_4": 0.35362373458014595, "updates_nfe_4": 9, "loss_nfe_1": 0.4059716686606407, "updates_nfe_1": 12, "loss_nfe_5": 0.26390672425429024, "updates_nfe_5": 15, "loss_nfe_3": 0.30862311167376383, "updates_nfe_3": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.485892998826298, "grad_norm": 0.01873681519408226, "grad_rms": 5.455977571227022e-05, "grad_to_weight": 0.0011365362613609234, "update_to_weight": 5.00883456272451e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.247841090061492, "grad_norm": 0.012954448499573834, "grad_rms": 3.772208869499515e-05, "grad_to_weight": 0.000797302757195098, "update_to_weight": 6.324997990041444e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06016656292372, "grad_norm": 0.000867312361732782, "grad_rms": 1.5527382066877856e-05, "grad_to_weight": 0.00017139996301458832, "update_to_weight": 2.8189742628336495e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.843419722791172, "grad_norm": 0.008789574505542063, "grad_rms": 2.81266744199009e-05, "grad_to_weight": 0.0006843640319520668, "update_to_weight": 7.292745705883703e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.055899190344365, "grad_norm": 0.02416139750965697, "grad_rms": 0.0005496908267190922, "grad_to_weight": 0.004778852702561757, "update_to_weight": 3.230738960534516e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636645834396447, "grad_norm": 0.07662291138355434, "grad_rms": 0.00018463405743264784, "grad_to_weight": 0.004605670646972339, "update_to_weight": 9.506206319635412e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.559274398572427, "grad_norm": 0.02587982392033598, "grad_rms": 3.5705106726843405e-05, "grad_to_weight": 0.0010537698915828032, "update_to_weight": 0.00010345995111780078, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8254778228629, "grad_norm": 0.010099509889804834, "grad_rms": 1.2260569724371346e-05, "grad_to_weight": 0.00032763514479292657, "update_to_weight": 8.860670069704136e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.601713316729345, "grad_norm": 0.006437563007841212, "grad_rms": 5.779014716640603e-06, "grad_to_weight": 0.00018604752166213572, "update_to_weight": 0.00010318096658586067, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03428752046457, "grad_norm": 0.005083961303771157, "grad_rms": 4.089344342249016e-06, "grad_to_weight": 0.00013727714623811352, "update_to_weight": 0.00010442364843053298, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660655895956012, "grad_norm": 0.0017431674367540757, "grad_rms": 3.1493655608226607e-06, "grad_to_weight": 7.0686174938272e-05, "update_to_weight": 6.570410870543633e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.28974586242678, "grad_norm": 0.005160638827913053, "grad_rms": 4.151020814745912e-06, "grad_to_weight": 0.0001383929739546153, "update_to_weight": 0.00010695729475614236, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.40139058510468, "grad_norm": 0.013298183490666719, "grad_rms": 5.903932713186806e-06, "grad_to_weight": 0.0002097459277776622, "update_to_weight": 0.00012091607114330958, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.754307756675935, "grad_norm": 0.02480285301708813, "grad_rms": 1.850716982766327e-05, "grad_to_weight": 0.0004886846873372183, "update_to_weight": 9.314267823001392e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0664883203922, "grad_norm": 0.03417152579817678, "grad_rms": 2.934515770956325e-05, "grad_to_weight": 0.0008528704967833399, "update_to_weight": 0.00011283516357405152, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.196688150494733, "grad_norm": 0.060941242445442016, "grad_rms": 9.499876106579068e-05, "grad_to_weight": 0.0021612907913219855, "update_to_weight": 9.184497962235395e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406920679274736, "grad_norm": 0.006301373583796572, "grad_rms": 0.0005364085415957492, "grad_to_weight": 0.0007495459781523942, "update_to_weight": 7.2061661614608606e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3088365016503658, "grad_norm": 0.0716221910106036, "grad_rms": 0.0016580277034998815, "grad_to_weight": 0.05472203053650493, "update_to_weight": 0.00012990767279180864, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42750}
{"step": 42800, "lr": 1.5007902737376136e-05, "window_steps": 50, "elapsed_seconds": 1041.012631399, "steps_per_second": 1.7290856476746197, "eta_seconds": 4164.050525596, "images_seen": 43827200, "images_per_second": 1770.5837032188106, "effective_batch_size": 1024, "loss": 0.3756271302700043, "edge_loss": 0.19372396171092987, "diag_loss": 0.18190322816371918, "jvp_rms": 2.3044140338897705, "grad_norm": 0.24846240878105164, "grad_norm_postclip": 0.24846240878105164, "clip_fraction": 0.0, "ema_updates": 42800, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39672968089580535, "updates_nfe_1": 10, "loss_nfe_4": 0.315699602663517, "updates_nfe_4": 10, "loss_nfe_3": 0.3588660180568695, "updates_nfe_3": 10, "loss_nfe_5": 0.4477489482272755, "updates_nfe_5": 11, "loss_nfe_2": 0.34924061761962044, "updates_nfe_2": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.486618626934998, "grad_norm": 0.049056375597367634, "grad_rms": 0.00014284737412015452, "grad_to_weight": 0.0029755268019132693, "update_to_weight": 5.247289193707247e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2464454871457, "grad_norm": 0.019415292137261463, "grad_rms": 5.653543430004854e-05, "grad_to_weight": 0.0011950486124871424, "update_to_weight": 7.601451722294624e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.059819206355201, "grad_norm": 0.0023305866720033574, "grad_rms": 4.172419452649164e-05, "grad_to_weight": 0.000460606708847642, "update_to_weight": 4.093181274382064e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.843523083837944, "grad_norm": 0.010268123100555438, "grad_rms": 3.285803598009037e-05, "grad_to_weight": 0.0007994786970466583, "update_to_weight": 8.165978173092409e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.055570027164152, "grad_norm": 0.030760485003723047, "grad_rms": 0.0006998252657040458, "grad_to_weight": 0.006084474122293523, "update_to_weight": 4.067936821134582e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636819180402433, "grad_norm": 0.06875742078417502, "grad_rms": 0.00016568101301239302, "grad_to_weight": 0.004132846551891888, "update_to_weight": 9.25133692242165e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.558469174668137, "grad_norm": 0.026103629712455945, "grad_rms": 3.60138804541426e-05, "grad_to_weight": 0.001062917624335544, "update_to_weight": 0.00010218507903627739, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.825042214716632, "grad_norm": 0.014390799696723444, "grad_rms": 1.7470095578524096e-05, "grad_to_weight": 0.0004668541764349287, "update_to_weight": 7.675061681790462e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60123881509248, "grad_norm": 0.010378852263418073, "grad_rms": 9.317118900284808e-06, "grad_to_weight": 0.00029995608882335713, "update_to_weight": 8.500523274666435e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.035418080302705, "grad_norm": 0.008649303878954753, "grad_rms": 6.957169767511705e-06, "grad_to_weight": 0.00023354141325475917, "update_to_weight": 8.76340250259166e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.660611345891752, "grad_norm": 0.0034274751102499807, "grad_rms": 6.192389695448975e-06, "grad_to_weight": 0.00013898581272685963, "update_to_weight": 5.784461961567386e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.29045294521609, "grad_norm": 0.008465993316758113, "grad_rms": 6.809721750974429e-06, "grad_to_weight": 0.0002270284388659912, "update_to_weight": 9.431335443838034e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.402049053786506, "grad_norm": 0.019535174571560555, "grad_rms": 8.672940653270306e-06, "grad_to_weight": 0.0003081158237486628, "update_to_weight": 0.0001088457164099041, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.754767132937864, "grad_norm": 0.027158537793695092, "grad_rms": 2.0264913511064107e-05, "grad_to_weight": 0.0005350933385737132, "update_to_weight": 8.974378382508893e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.06701082493318, "grad_norm": 0.027126151052099784, "grad_rms": 2.3294867937087402e-05, "grad_to_weight": 0.0006770195852798565, "update_to_weight": 0.00010500062969961492, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.197222806926174, "grad_norm": 0.0423235571316214, "grad_rms": 6.597642795026346e-05, "grad_to_weight": 0.0015009831791386679, "update_to_weight": 9.306077383220599e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406981936216997, "grad_norm": 0.0030860420501516157, "grad_rms": 0.00026270134493876804, "grad_to_weight": 0.0003670808470346593, "update_to_weight": 5.721072185616476e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.308987595233468, "grad_norm": 0.05551143266656349, "grad_rms": 0.0012850695004360338, "grad_to_weight": 0.04240791346587406, "update_to_weight": 0.00016023685119040767, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42800}
{"step": 42850, "lr": 1.493989654492181e-05, "window_steps": 50, "elapsed_seconds": 1068.483481775, "steps_per_second": 1.7314259242704613, "eta_seconds": 4129.544267400675, "images_seen": 43878400, "images_per_second": 1772.9801464529523, "effective_batch_size": 1024, "loss": 0.35607630014419556, "edge_loss": 0.18464350700378418, "diag_loss": 0.17143283784389496, "jvp_rms": 1.918984293937683, "grad_norm": 0.31351667642593384, "grad_norm_postclip": 0.28892096877098083, "clip_fraction": 0.019999999552965164, "ema_updates": 42850, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.29385847846666974, "updates_nfe_4": 12, "loss_nfe_3": 0.33663914524591887, "updates_nfe_3": 13, "loss_nfe_1": 0.4054825693368912, "updates_nfe_1": 10, "loss_nfe_2": 0.38274507224559784, "updates_nfe_2": 8, "loss_nfe_5": 0.3977743685245514, "updates_nfe_5": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.484212084503007, "grad_norm": 0.12443397563228469, "grad_rms": 0.00036233958285652415, "grad_to_weight": 0.00754867596912724, "update_to_weight": 4.831423880845356e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248651536598736, "grad_norm": 0.03734505470293116, "grad_rms": 0.00010874515158787027, "grad_to_weight": 0.0022983479348310555, "update_to_weight": 5.479304780851054e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.059862368072662, "grad_norm": 0.004571694252197568, "grad_rms": 8.184645633898151e-05, "grad_to_weight": 0.0009035214635569147, "update_to_weight": 3.445456226462005e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.843034784678766, "grad_norm": 0.027251111144564175, "grad_rms": 8.720366728337092e-05, "grad_to_weight": 0.0021218591712508384, "update_to_weight": 7.372628196763832e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.054999741309939, "grad_norm": 0.03593776387915965, "grad_rms": 0.0008176124385716985, "grad_to_weight": 0.007109350290460516, "update_to_weight": 2.9962721221895962e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636854949988386, "grad_norm": 0.12314239079598906, "grad_rms": 0.00029672951398059406, "grad_to_weight": 0.007401783039292232, "update_to_weight": 7.638367299257984e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55871646074401, "grad_norm": 0.03732942758172066, "grad_rms": 5.1501555805021284e-05, "grad_to_weight": 0.0015200072707948747, "update_to_weight": 9.77708338806101e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8247848070094, "grad_norm": 0.026323587078877064, "grad_rms": 3.195622146990834e-05, "grad_to_weight": 0.0008539747233820498, "update_to_weight": 8.663417335388542e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60000049392612, "grad_norm": 0.018221923618165107, "grad_rms": 1.6357861604866883e-05, "grad_to_weight": 0.0005266451837584189, "update_to_weight": 0.00010342739353401811, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.035752627106824, "grad_norm": 0.013591358972054867, "grad_rms": 1.0932370172570015e-05, "grad_to_weight": 0.0003669794187497413, "update_to_weight": 0.00010182766406580424, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.661462733192145, "grad_norm": 0.004757658941202645, "grad_rms": 8.595621340577788e-06, "grad_to_weight": 0.00019291876530905271, "update_to_weight": 6.530588548394234e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.291006162679494, "grad_norm": 0.01167025503277563, "grad_rms": 9.387107520956909e-06, "grad_to_weight": 0.00031295092929015995, "update_to_weight": 0.00010214151054520326, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.40275564853727, "grad_norm": 0.021763544965808942, "grad_rms": 9.662259899536806e-06, "grad_to_weight": 0.0003432586603404364, "update_to_weight": 0.00011214952289395777, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75501726276674, "grad_norm": 0.029406444972480877, "grad_rms": 2.194223667569963e-05, "grad_to_weight": 0.0005793800605019808, "update_to_weight": 8.42721517860408e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.06765975694659, "grad_norm": 0.03482906595311443, "grad_rms": 2.9909827243519076e-05, "grad_to_weight": 0.0008692563070663506, "update_to_weight": 9.337644328966831e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.197161116316856, "grad_norm": 0.07496058998206263, "grad_rms": 0.00011685293721131353, "grad_to_weight": 0.002658444574361253, "update_to_weight": 7.253001066347385e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406795441470344, "grad_norm": 0.00496720789352235, "grad_rms": 0.00042283681589973075, "grad_to_weight": 0.0005908562814576572, "update_to_weight": 7.726882035811924e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3087876816710837, "grad_norm": 0.06791149281812155, "grad_rms": 0.0015721263883397333, "grad_to_weight": 0.0518888539135782, "update_to_weight": 0.00021542725958916797, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42850}
{"step": 42900, "lr": 1.4872328443340663e-05, "window_steps": 50, "elapsed_seconds": 1095.9865722060001, "steps_per_second": 1.7335978817474769, "eta_seconds": 4095.5287698224215, "images_seen": 43929600, "images_per_second": 1775.2042309094163, "effective_batch_size": 1024, "loss": 0.3624397814273834, "edge_loss": 0.18647556006908417, "diag_loss": 0.17596417665481567, "jvp_rms": 1.7647265195846558, "grad_norm": 0.19453175365924835, "grad_norm_postclip": 0.18561773002147675, "clip_fraction": 0.03999999910593033, "ema_updates": 42900, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3439585119485855, "updates_nfe_4": 12, "loss_nfe_3": 0.3545133024454117, "updates_nfe_3": 15, "loss_nfe_5": 0.3887891359627247, "updates_nfe_5": 8, "loss_nfe_1": 0.39854029193520546, "updates_nfe_1": 8, "loss_nfe_2": 0.3397361934185028, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.488953562904147, "grad_norm": 0.012521013789169458, "grad_rms": 3.645996915436776e-05, "grad_to_weight": 0.0007593576961329116, "update_to_weight": 7.149985076894015e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24551564476958, "grad_norm": 0.01204629755634627, "grad_rms": 3.507764185266137e-05, "grad_to_weight": 0.0007415152476384896, "update_to_weight": 0.0001129310158709608, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060270030238571, "grad_norm": 0.0032793203184986446, "grad_rms": 5.8709251420415245e-05, "grad_to_weight": 0.0006480524357203203, "update_to_weight": 4.003511605651424e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.843273590236715, "grad_norm": 0.009919096833966795, "grad_rms": 3.174115049739238e-05, "grad_to_weight": 0.0007723184252266619, "update_to_weight": 8.388475122639844e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.055368368761279, "grad_norm": 0.043985205662812873, "grad_rms": 0.001000698078599849, "grad_to_weight": 0.008700692502372603, "update_to_weight": 2.2279257361641528e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.636938412056544, "grad_norm": 0.07338728202270249, "grad_rms": 0.00017683733754227698, "grad_to_weight": 0.004411104988494747, "update_to_weight": 6.506873205274196e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.558618292197853, "grad_norm": 0.02476488772657662, "grad_rms": 3.416688467732949e-05, "grad_to_weight": 0.0010083990651234764, "update_to_weight": 8.649420879768e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.824788767144252, "grad_norm": 0.014339419113988528, "grad_rms": 1.740772074806463e-05, "grad_to_weight": 0.00046519115580356325, "update_to_weight": 6.85673061867921e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60183855585296, "grad_norm": 0.010416913968679737, "grad_rms": 9.351286978263955e-06, "grad_to_weight": 0.0003010508806306681, "update_to_weight": 8.056912257642632e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0363294248448, "grad_norm": 0.008672343620131742, "grad_rms": 6.975702055544486e-06, "grad_to_weight": 0.00023415775145131253, "update_to_weight": 8.512131001620294e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.661441696360104, "grad_norm": 0.003412533354086886, "grad_rms": 6.165394553567547e-06, "grad_to_weight": 0.00013837525786623242, "update_to_weight": 5.653706568147343e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.292187860068296, "grad_norm": 0.008706728332990103, "grad_rms": 7.003359805591264e-06, "grad_to_weight": 0.0002334732509033906, "update_to_weight": 9.149824318284467e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.404567344602626, "grad_norm": 0.02040661311652593, "grad_rms": 9.059829173553089e-06, "grad_to_weight": 0.00032184768339505213, "update_to_weight": 0.00010633355113536177, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75559929086642, "grad_norm": 0.028777878963891026, "grad_rms": 2.147321894371314e-05, "grad_to_weight": 0.0005669892458361666, "update_to_weight": 8.113258573118292e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.06806799986977, "grad_norm": 0.030708397337993142, "grad_rms": 2.6371159667076424e-05, "grad_to_weight": 0.000766405740803199, "update_to_weight": 9.974758059046888e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19728882655155, "grad_norm": 0.06567223983035242, "grad_rms": 0.00010237371556518977, "grad_to_weight": 0.00232902674559666, "update_to_weight": 7.452635306624207e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407314530934972, "grad_norm": 0.008182474505213625, "grad_rms": 0.0006965384860330059, "grad_to_weight": 0.0009732566178065491, "update_to_weight": 4.410887129711943e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3093369380544568, "grad_norm": 0.096224880687266, "grad_rms": 0.0022275710320260583, "grad_to_weight": 0.07349130532454581, "update_to_weight": 0.00012447476744692807, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42900}
{"step": 42950, "lr": 1.4805199171546641e-05, "window_steps": 50, "elapsed_seconds": 1123.461673245, "steps_per_second": 1.73570674143928, "eta_seconds": 4061.746049424231, "images_seen": 43980800, "images_per_second": 1777.3637032338227, "effective_batch_size": 1024, "loss": 0.35867637395858765, "edge_loss": 0.18720579147338867, "diag_loss": 0.17147065699100494, "jvp_rms": 1.8673436641693115, "grad_norm": 0.20824237167835236, "grad_norm_postclip": 0.20824237167835236, "clip_fraction": 0.0, "ema_updates": 42950, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.40951323780146515, "updates_nfe_1": 11, "loss_nfe_2": 0.37095671594142915, "updates_nfe_2": 10, "loss_nfe_4": 0.38374221622943877, "updates_nfe_4": 10, "loss_nfe_3": 0.324347123503685, "updates_nfe_3": 12, "loss_nfe_5": 0.2842887895447867, "updates_nfe_5": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48659086114095, "grad_norm": 0.08837553407758583, "grad_rms": 0.0002573409230853733, "grad_to_weight": 0.005360449277957628, "update_to_weight": 6.675220533008647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251169954878442, "grad_norm": 0.06642191716109584, "grad_rms": 0.00019341413496104482, "grad_to_weight": 0.004087208326878437, "update_to_weight": 7.438536478807742e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060713652995228, "grad_norm": 0.0031727974850509866, "grad_rms": 5.680218678399792e-05, "grad_to_weight": 0.0006269466527064099, "update_to_weight": 3.482772268895467e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.845779592719415, "grad_norm": 0.03213863768933639, "grad_rms": 0.00010284377224598917, "grad_to_weight": 0.002501883008140009, "update_to_weight": 8.118596326796788e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.05464052063514, "grad_norm": 0.01933690776702515, "grad_rms": 0.00043992988453579724, "grad_to_weight": 0.0038255752685248073, "update_to_weight": 3.989565482224119e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.6367494753473, "grad_norm": 0.09124175088318903, "grad_rms": 0.00021986027897705044, "grad_to_weight": 0.005484349633225713, "update_to_weight": 8.605502631270581e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.559460789231764, "grad_norm": 0.016826649096989457, "grad_rms": 2.321489140391926e-05, "grad_to_weight": 0.0006851391910186887, "update_to_weight": 0.00010075550008950406, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.825265959669757, "grad_norm": 0.00785580661072964, "grad_rms": 9.536766213700827e-06, "grad_to_weight": 0.0002548495971132183, "update_to_weight": 7.529237161007357e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60161100837647, "grad_norm": 0.003803838294578804, "grad_rms": 3.4147141484000238e-06, "grad_to_weight": 0.00010993240440908829, "update_to_weight": 8.489735289718131e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03679415248774, "grad_norm": 0.0031651337107779474, "grad_rms": 2.5459126966663565e-06, "grad_to_weight": 8.545917062223236e-05, "update_to_weight": 8.626636434397532e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66163968936309, "grad_norm": 0.001182125478430164, "grad_rms": 2.1357358976779684e-06, "grad_to_weight": 4.793377461191404e-05, "update_to_weight": 5.418048209325576e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.29308474666416, "grad_norm": 0.0032547994880922425, "grad_rms": 2.6180364240601077e-06, "grad_to_weight": 8.727622051654985e-05, "update_to_weight": 9.159374039617318e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.406159515090685, "grad_norm": 0.008390522241101914, "grad_rms": 3.725101159472681e-06, "grad_to_weight": 0.000132329765834579, "update_to_weight": 0.00010415071051884244, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75529865735079, "grad_norm": 0.021165096908688827, "grad_rms": 1.5792781686775544e-05, "grad_to_weight": 0.00041700270648734576, "update_to_weight": 8.586627637820268e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.068620949137056, "grad_norm": 0.04883610660167924, "grad_rms": 4.193852093732915e-05, "grad_to_weight": 0.0012188117645394283, "update_to_weight": 0.00011564965644851187, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.197328871167034, "grad_norm": 0.0928065040286, "grad_rms": 0.00014467218828787966, "grad_to_weight": 0.0032913225381251835, "update_to_weight": 9.57149580623894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406896630130944, "grad_norm": 0.003960435024382893, "grad_rms": 0.0003371346179151622, "grad_to_weight": 0.0004710935793106339, "update_to_weight": 5.867649211335235e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3090209263701817, "grad_norm": 0.05008638182591596, "grad_rms": 0.0011594815442485828, "grad_to_weight": 0.03826247603604153, "update_to_weight": 0.00018068361138796608, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 42950}
{"step": 43000, "lr": 1.4738509463654698e-05, "window_steps": 50, "elapsed_seconds": 1151.121386896, "steps_per_second": 1.7374362276362547, "eta_seconds": 4028.9248541360002, "images_seen": 44032000, "images_per_second": 1779.1346970995248, "effective_batch_size": 1024, "loss": 0.37084653973579407, "edge_loss": 0.19238755106925964, "diag_loss": 0.17845900356769562, "jvp_rms": 2.1513280868530273, "grad_norm": 0.2333534210920334, "grad_norm_postclip": 0.2333534210920334, "clip_fraction": 0.0, "ema_updates": 43000, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3449607938528061, "updates_nfe_5": 13, "loss_nfe_1": 0.40004437913497287, "updates_nfe_1": 12, "loss_nfe_2": 0.3548620730638504, "updates_nfe_2": 10, "loss_nfe_3": 0.3384616772333781, "updates_nfe_3": 6, "loss_nfe_4": 0.4086575077639686, "updates_nfe_4": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.488359448849152, "grad_norm": 0.03219175626566559, "grad_rms": 9.373924989095845e-05, "grad_to_weight": 0.0019523929209289827, "update_to_weight": 4.9606951719810917e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24716866109148, "grad_norm": 0.03535911850766698, "grad_rms": 0.00010296229936511288, "grad_to_weight": 0.0021763249489952403, "update_to_weight": 6.639395846414536e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060731932283085, "grad_norm": 0.0033985068412895475, "grad_rms": 6.084303246430538e-05, "grad_to_weight": 0.0006715445288872186, "update_to_weight": 4.3297111187951426e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.844351724233508, "grad_norm": 0.014130306537078264, "grad_rms": 4.521703879649715e-05, "grad_to_weight": 0.0011001183119595315, "update_to_weight": 7.765805504250839e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.055116897574628, "grad_norm": 0.01773729710935013, "grad_rms": 0.0004035374819649355, "grad_to_weight": 0.0035087807994826446, "update_to_weight": 3.244397383229698e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63725483038039, "grad_norm": 0.060005738485615766, "grad_rms": 0.00014459256070788702, "grad_to_weight": 0.0036067091054013635, "update_to_weight": 0.00010374202807685612, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55906439673484, "grad_norm": 0.025914607298983466, "grad_rms": 3.57530956254832e-05, "grad_to_weight": 0.0010551952175518887, "update_to_weight": 0.00011914955131137704, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.825027364330413, "grad_norm": 0.016427547214470522, "grad_rms": 1.9942659616258888e-05, "grad_to_weight": 0.0005329288769255036, "update_to_weight": 0.00010236779198745597, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60227600726306, "grad_norm": 0.011116199321903292, "grad_rms": 9.979037004553068e-06, "grad_to_weight": 0.00032125630463065465, "update_to_weight": 0.00012035266552852559, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03753243122104, "grad_norm": 0.009065523112865691, "grad_rms": 7.29196062598389e-06, "grad_to_weight": 0.0002447658501467513, "update_to_weight": 0.00011858315787592631, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.662082692950047, "grad_norm": 0.003556068407511744, "grad_rms": 6.424718095584161e-06, "grad_to_weight": 0.00014419173156565113, "update_to_weight": 7.911415459151151e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.293817951798324, "grad_norm": 0.00876142997609683, "grad_rms": 7.047359718530076e-06, "grad_to_weight": 0.0002349298210073541, "update_to_weight": 0.00011687493411941078, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.40797688937193, "grad_norm": 0.02015351246732352, "grad_rms": 8.947461254761497e-06, "grad_to_weight": 0.00031783875556360054, "update_to_weight": 0.000125554000106632, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7556882780455, "grad_norm": 0.02597363583670553, "grad_rms": 1.9380774023883817e-05, "grad_to_weight": 0.000511738422192582, "update_to_weight": 9.601184654771432e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0692211113312, "grad_norm": 0.029642003420468882, "grad_rms": 2.5455382658022354e-05, "grad_to_weight": 0.0007397698931583774, "update_to_weight": 0.00010643477157151878, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.197339694026322, "grad_norm": 0.06308693146882169, "grad_rms": 9.834358619035262e-05, "grad_to_weight": 0.0022373362931889214, "update_to_weight": 8.44753101623766e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407417528268054, "grad_norm": 0.009870091857068063, "grad_rms": 0.0008401980152517934, "grad_to_weight": 0.0011739742702062904, "update_to_weight": 4.390561094468302e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3095573407582524, "grad_norm": 0.12110525891772134, "grad_rms": 0.0028035427496958362, "grad_to_weight": 0.09247801157572808, "update_to_weight": 0.00016403383137561238, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43000}
{"step": 43050, "lr": 1.4672260048972808e-05, "window_steps": 50, "elapsed_seconds": 1202.066354538, "steps_per_second": 1.7053967048165932, "eta_seconds": 4075.2981287995613, "images_seen": 44083200, "images_per_second": 1746.3262257321915, "effective_batch_size": 1024, "loss": 0.33856603503227234, "edge_loss": 0.1777298003435135, "diag_loss": 0.16083618998527527, "jvp_rms": 1.4663281440734863, "grad_norm": 0.27545803785324097, "grad_norm_postclip": 0.27545803785324097, "clip_fraction": 0.0, "ema_updates": 43050, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4052434219763829, "updates_nfe_1": 13, "loss_nfe_2": 0.38480775952339175, "updates_nfe_2": 5, "loss_nfe_4": 0.301139418895428, "updates_nfe_4": 13, "loss_nfe_5": 0.3027322980073782, "updates_nfe_5": 13, "loss_nfe_3": 0.3142941892147064, "updates_nfe_3": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48919971603809, "grad_norm": 0.05992654690192136, "grad_rms": 0.00017450025120663613, "grad_to_weight": 0.00363429080452184, "update_to_weight": 7.57239177122169e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.247161147761418, "grad_norm": 0.017163619810458757, "grad_rms": 4.9978784418232885e-05, "grad_to_weight": 0.001056407310444115, "update_to_weight": 9.262283851039683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060823893059558, "grad_norm": 0.0037471803230008692, "grad_rms": 6.708528912522022e-05, "grad_to_weight": 0.0007404289108221634, "update_to_weight": 4.5260575710383456e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.845552712410031, "grad_norm": 0.016859593055360552, "grad_rms": 5.3950766834179506e-05, "grad_to_weight": 0.001312484829015771, "update_to_weight": 0.0001424514404474746, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0551663250071535, "grad_norm": 0.0470619843266011, "grad_rms": 0.001070697216963159, "grad_to_weight": 0.009309680691175775, "update_to_weight": 4.4086255257264454e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.637515297760153, "grad_norm": 0.08333599345858969, "grad_rms": 0.00020081020578059723, "grad_to_weight": 0.005008920621086306, "update_to_weight": 0.00016527870073750246, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55926942813548, "grad_norm": 0.03473843377782378, "grad_rms": 4.7926890437069e-05, "grad_to_weight": 0.0014144734182535124, "update_to_weight": 0.00021371338016322382, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.825297640239825, "grad_norm": 0.02061336959647706, "grad_rms": 2.5024150473575833e-05, "grad_to_weight": 0.0006687159954481039, "update_to_weight": 0.00016272948330340662, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60177858224469, "grad_norm": 0.013265903039099105, "grad_rms": 1.1908830841593528e-05, "grad_to_weight": 0.0003833878945721673, "update_to_weight": 0.00018756898804832883, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0379131004426, "grad_norm": 0.009421241175408872, "grad_rms": 7.57808665243844e-06, "grad_to_weight": 0.0002543674950005833, "update_to_weight": 0.00017844674225138808, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.662957539787783, "grad_norm": 0.0038946494501653924, "grad_rms": 7.03642954268783e-06, "grad_to_weight": 0.00015791493959644974, "update_to_weight": 0.00012272561320494368, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.294109265908205, "grad_norm": 0.00970043312276853, "grad_rms": 7.802657993980676e-06, "grad_to_weight": 0.00026010630937996475, "update_to_weight": 0.00017360283502779628, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.40927443193444, "grad_norm": 0.022680836215972444, "grad_rms": 1.0069505432218866e-05, "grad_to_weight": 0.0003576895717410989, "update_to_weight": 0.00018596424776018338, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.755875872128094, "grad_norm": 0.03688349176984104, "grad_rms": 2.752139221852339e-05, "grad_to_weight": 0.0007266841747104026, "update_to_weight": 0.00014014229047394844, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.06967198869052, "grad_norm": 0.058448101856747564, "grad_rms": 5.0192923106242646e-05, "grad_to_weight": 0.0014586618496214361, "update_to_weight": 0.00016703917297950356, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19740679566119, "grad_norm": 0.19637330937522862, "grad_rms": 0.000306118161501829, "grad_to_weight": 0.006964232945188602, "update_to_weight": 0.00012687328146305808, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407517801992471, "grad_norm": 0.03174027840126124, "grad_rms": 0.0027019119277174484, "grad_to_weight": 0.0037752258334486327, "update_to_weight": 1.0504979269072937e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3096664814410455, "grad_norm": 0.3907460215970678, "grad_rms": 0.009045628452561374, "grad_to_weight": 0.2983553653805999, "update_to_weight": 0.0003880810072541255, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43050}
{"step": 43100, "lr": 1.4606451651993947e-05, "window_steps": 50, "elapsed_seconds": 1229.582142704, "steps_per_second": 1.7078972823904597, "eta_seconds": 4040.0556117417145, "images_seen": 44134400, "images_per_second": 1748.8868171678307, "effective_batch_size": 1024, "loss": 0.34965747594833374, "edge_loss": 0.18121448159217834, "diag_loss": 0.1684430092573166, "jvp_rms": 1.7147655487060547, "grad_norm": 0.3617490828037262, "grad_norm_postclip": 0.3587269186973572, "clip_fraction": 0.019999999552965164, "ema_updates": 43100, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.35004792114098865, "updates_nfe_4": 12, "loss_nfe_5": 0.3147326409816742, "updates_nfe_5": 10, "loss_nfe_1": 0.4170072575410207, "updates_nfe_1": 6, "loss_nfe_2": 0.3630124628543854, "updates_nfe_2": 13, "loss_nfe_3": 0.32375186681747437, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.485262676040286, "grad_norm": 0.02718271901478375, "grad_rms": 7.915342267797324e-05, "grad_to_weight": 0.0016489102751326597, "update_to_weight": 7.56198103715007e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251646926758877, "grad_norm": 0.00697504642736137, "grad_rms": 2.0310653903428645e-05, "grad_to_weight": 0.0004291901281633632, "update_to_weight": 7.928945064242342e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061544069060064, "grad_norm": 0.0010951719476553715, "grad_rms": 1.960672303366867e-05, "grad_to_weight": 0.0002163711177286552, "update_to_weight": 3.373882459742626e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84652375576451, "grad_norm": 0.00785972757735394, "grad_rms": 2.515116044103857e-05, "grad_to_weight": 0.000611817463368416, "update_to_weight": 8.762570414694048e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.054515994904856, "grad_norm": 0.013730123555068557, "grad_rms": 0.000312371126915335, "grad_to_weight": 0.0027164071829842945, "update_to_weight": 3.692952770153262e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63747310965849, "grad_norm": 0.04782881244609661, "grad_rms": 0.00011525048506579344, "grad_to_weight": 0.0028747642223584256, "update_to_weight": 8.775895725801872e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.559988888418065, "grad_norm": 0.019292333505025405, "grad_rms": 2.661667362680597e-05, "grad_to_weight": 0.0007855188205774488, "update_to_weight": 9.282083650402745e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826292591575235, "grad_norm": 0.010680203736279596, "grad_rms": 1.2965518525936883e-05, "grad_to_weight": 0.00034646410055805656, "update_to_weight": 7.899069643338549e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60427798246194, "grad_norm": 0.007791777598321866, "grad_rms": 6.9946961846658025e-06, "grad_to_weight": 0.00022516804431726265, "update_to_weight": 9.183218622426866e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.040002591687106, "grad_norm": 0.006584900100279322, "grad_rms": 5.296642196976942e-06, "grad_to_weight": 0.0001777780680219815, "update_to_weight": 9.795910647365974e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66256033594685, "grad_norm": 0.002335440433416336, "grad_rms": 4.2194200713444225e-06, "grad_to_weight": 9.469578184923165e-05, "update_to_weight": 5.986058991860155e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.2958014612989, "grad_norm": 0.006557769723909032, "grad_rms": 5.274819558195126e-06, "grad_to_weight": 0.00017583131255977693, "update_to_weight": 9.623834382969029e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.412029215431595, "grad_norm": 0.016042613087130484, "grad_rms": 7.122364364770833e-06, "grad_to_weight": 0.00025299006017657046, "update_to_weight": 0.000105836033984874, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75657573624082, "grad_norm": 0.02313917883818882, "grad_rms": 1.7265784389239117e-05, "grad_to_weight": 0.00045588534101340415, "update_to_weight": 8.490298184350383e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07043664072211, "grad_norm": 0.022852707389940934, "grad_rms": 1.962500317981408e-05, "grad_to_weight": 0.0005703134107282616, "update_to_weight": 0.00010325469764689319, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.197783427680534, "grad_norm": 0.0332977999753838, "grad_rms": 5.190655157235911e-05, "grad_to_weight": 0.0011808658670204835, "update_to_weight": 8.330477132847647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.406997363820985, "grad_norm": 0.002593003462499002, "grad_rms": 0.00022073111317323237, "grad_to_weight": 0.0003084339568914151, "update_to_weight": 4.032131360421601e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.309551423780748, "grad_norm": 0.03702281594303065, "grad_rms": 0.0008570647396982581, "grad_to_weight": 0.028271372372795957, "update_to_weight": 9.526568960216653e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43100}
{"step": 43150, "lr": 1.4541084992388194e-05, "window_steps": 50, "elapsed_seconds": 1257.0749694620001, "steps_per_second": 1.7103196326629204, "eta_seconds": 4005.0993213091633, "images_seen": 44185600, "images_per_second": 1751.3673038468305, "effective_batch_size": 1024, "loss": 0.3704785108566284, "edge_loss": 0.19210760295391083, "diag_loss": 0.17837077379226685, "jvp_rms": 2.171211004257202, "grad_norm": 0.23501534759998322, "grad_norm_postclip": 0.23501534759998322, "clip_fraction": 0.0, "ema_updates": 43150, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.37832407309458804, "updates_nfe_2": 13, "loss_nfe_1": 0.4024971753358841, "updates_nfe_1": 10, "loss_nfe_3": 0.34921538829803467, "updates_nfe_3": 9, "loss_nfe_5": 0.3712481204420328, "updates_nfe_5": 8, "loss_nfe_4": 0.3467812955379486, "updates_nfe_4": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.489272820960373, "grad_norm": 0.04082878306674599, "grad_rms": 0.00011888942830744011, "grad_to_weight": 0.0024760814809763106, "update_to_weight": 6.773869095207051e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24667653063283, "grad_norm": 0.012624825752364188, "grad_rms": 3.676225945127711e-05, "grad_to_weight": 0.0007770712815362758, "update_to_weight": 0.0001250243893577976, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060898327300447, "grad_norm": 0.0013564085781003995, "grad_rms": 2.4283608951309348e-05, "grad_to_weight": 0.00026801735391193416, "update_to_weight": 3.592255346565736e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.845356713105092, "grad_norm": 0.011241773042045822, "grad_rms": 3.597371978093742e-05, "grad_to_weight": 0.000875162386932917, "update_to_weight": 9.04192336163785e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.055117840850764, "grad_norm": 0.030774272900098633, "grad_rms": 0.000700138951208139, "grad_to_weight": 0.00608774589810935, "update_to_weight": 4.050400744623919e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63815452612327, "grad_norm": 0.05834113342780414, "grad_rms": 0.0001405814525380579, "grad_to_weight": 0.0035064666178093113, "update_to_weight": 8.177978817799366e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.559886997223874, "grad_norm": 0.019463116299497513, "grad_rms": 2.6852294159716306e-05, "grad_to_weight": 0.0007924758082843592, "update_to_weight": 0.00010328918914538632, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8260193544071, "grad_norm": 0.011309006374412905, "grad_rms": 1.3728870279816007e-05, "grad_to_weight": 0.0003668656093540047, "update_to_weight": 8.502949353969974e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60323379487097, "grad_norm": 0.007985862135236482, "grad_rms": 7.168926307731901e-06, "grad_to_weight": 0.0002307836944540189, "update_to_weight": 0.00010079999789905552, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0393434589112, "grad_norm": 0.006659742034654355, "grad_rms": 5.356842191157167e-06, "grad_to_weight": 0.0001798018380655747, "update_to_weight": 0.00010516712755097117, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.662987237014192, "grad_norm": 0.002393718544702131, "grad_rms": 4.3247106319431445e-06, "grad_to_weight": 9.705712133320331e-05, "update_to_weight": 6.954799329216344e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.29584892021748, "grad_norm": 0.006968673835923403, "grad_rms": 5.605335135571058e-06, "grad_to_weight": 0.00018684851096513846, "update_to_weight": 0.00011005547813012058, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.411521004418034, "grad_norm": 0.01695581570545707, "grad_rms": 7.5277946865805875e-06, "grad_to_weight": 0.00026739329757246666, "update_to_weight": 0.0001253354196164905, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75615966698648, "grad_norm": 0.025083350351336492, "grad_rms": 1.871646880619468e-05, "grad_to_weight": 0.0004941932273030411, "update_to_weight": 9.22741379190404e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0704351175254, "grad_norm": 0.02652440398355696, "grad_rms": 2.277811130374448e-05, "grad_to_weight": 0.0006619444961294198, "update_to_weight": 0.0001150471519766473, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198073474197667, "grad_norm": 0.06447592967638145, "grad_rms": 0.00010050883756275307, "grad_to_weight": 0.0022865366932027973, "update_to_weight": 9.132024748364394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40726461987916, "grad_norm": 0.008941423768045732, "grad_rms": 0.0007611445376830428, "grad_to_weight": 0.001063535427076191, "update_to_weight": 4.715673675452624e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3096807719162282, "grad_norm": 0.09377759288990395, "grad_rms": 0.002170917208550276, "grad_to_weight": 0.07160339748494246, "update_to_weight": 0.0001732511109904041, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43150}
{"step": 43200, "lr": 1.4476160784994871e-05, "window_steps": 50, "elapsed_seconds": 1284.546477258, "steps_per_second": 1.7126667185263176, "eta_seconds": 3970.416384252, "images_seen": 44236800, "images_per_second": 1753.7707197709492, "effective_batch_size": 1024, "loss": 0.3757188022136688, "edge_loss": 0.19472026824951172, "diag_loss": 0.1809985637664795, "jvp_rms": 2.0646092891693115, "grad_norm": 0.29899388551712036, "grad_norm_postclip": 0.2971493899822235, "clip_fraction": 0.019999999552965164, "ema_updates": 43200, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.353925245148795, "updates_nfe_3": 7, "loss_nfe_2": 0.35967204968134564, "updates_nfe_2": 9, "loss_nfe_5": 0.3684745542705059, "updates_nfe_5": 8, "loss_nfe_1": 0.4012128439816562, "updates_nfe_1": 11, "loss_nfe_4": 0.3806852380434672, "updates_nfe_4": 15, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48496185270552, "grad_norm": 0.008202854326584979, "grad_rms": 2.388591057888282e-05, "grad_to_weight": 0.0004975961970599721, "update_to_weight": 9.990135219035475e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.249865722101465, "grad_norm": 0.007693100788954181, "grad_rms": 2.2401558067872483e-05, "grad_to_weight": 0.00047342549904832656, "update_to_weight": 9.992779163363157e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060768302216158, "grad_norm": 0.000619809186911438, "grad_rms": 1.1096364445339192e-05, "grad_to_weight": 0.00012247333801865967, "update_to_weight": 3.9309435794577406e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.844412904881786, "grad_norm": 0.007398378910101586, "grad_rms": 2.3674842816143272e-05, "grad_to_weight": 0.0005759997724216479, "update_to_weight": 9.527560107136179e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.054802776832237, "grad_norm": 0.025116784880838797, "grad_rms": 0.0005714266420291138, "grad_to_weight": 0.00496889512602885, "update_to_weight": 3.661799677751912e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.637518966285676, "grad_norm": 0.07300412908623452, "grad_rms": 0.0001759140747740014, "grad_to_weight": 0.004387921614645208, "update_to_weight": 9.361456985313742e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.55967327285791, "grad_norm": 0.013551384923907743, "grad_rms": 1.869617221871659e-05, "grad_to_weight": 0.0005517738274997346, "update_to_weight": 7.956383068256002e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826115383687863, "grad_norm": 0.007316736780696107, "grad_rms": 8.882347998406732e-06, "grad_to_weight": 0.0002373551350738107, "update_to_weight": 6.0981322071905256e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60407867206101, "grad_norm": 0.004475536971970942, "grad_rms": 4.01769955406814e-06, "grad_to_weight": 0.00012933553337411772, "update_to_weight": 6.63458889017414e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.04039641792289, "grad_norm": 0.003763177101161331, "grad_rms": 3.026955963669514e-06, "grad_to_weight": 0.0001015965665891315, "update_to_weight": 7.209873319725814e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663347313039115, "grad_norm": 0.0012667711981951193, "grad_rms": 2.288664588908681e-06, "grad_to_weight": 5.1362500885084556e-05, "update_to_weight": 4.575846262762546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.297703043803516, "grad_norm": 0.003913030894923264, "grad_rms": 3.1474926332209416e-06, "grad_to_weight": 0.00010491345513496328, "update_to_weight": 7.746881233226683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.413469124518414, "grad_norm": 0.010326356609460494, "grad_rms": 4.5845445460589e-06, "grad_to_weight": 0.00016284169202576988, "update_to_weight": 8.91851724757347e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75698458709895, "grad_norm": 0.022241617841798452, "grad_rms": 1.659605039616023e-05, "grad_to_weight": 0.0004381981715961052, "update_to_weight": 7.586114435331175e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.070803729441685, "grad_norm": 0.0354839391366653, "grad_rms": 3.0472206487705168e-05, "grad_to_weight": 0.0008855310059726548, "update_to_weight": 0.00011615322192922518, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.197758535489683, "grad_norm": 0.09514811696371996, "grad_rms": 0.00014832243102671412, "grad_to_weight": 0.003374314906767025, "update_to_weight": 0.0001070596978245752, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407350829698546, "grad_norm": 0.015361989552559479, "grad_rms": 0.0013076993932063922, "grad_to_weight": 0.001827209291456474, "update_to_weight": 4.646761862182392e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3099702797495862, "grad_norm": 0.18702289358670499, "grad_rms": 0.004329512046197508, "grad_to_weight": 0.14276880664991595, "update_to_weight": 0.00013358793324957414, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43200}
{"step": 43250, "lr": 1.4411679739814713e-05, "window_steps": 50, "elapsed_seconds": 1312.013020464, "steps_per_second": 1.7149220052742131, "eta_seconds": 3936.039061392, "images_seen": 44288000, "images_per_second": 1756.0801334007942, "effective_batch_size": 1024, "loss": 0.35525619983673096, "edge_loss": 0.1836819350719452, "diag_loss": 0.171574205160141, "jvp_rms": 2.0351953506469727, "grad_norm": 0.2273246943950653, "grad_norm_postclip": 0.2273246943950653, "clip_fraction": 0.0, "ema_updates": 43250, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.40140238404273987, "updates_nfe_1": 9, "loss_nfe_5": 0.2921165997783343, "updates_nfe_5": 12, "loss_nfe_3": 0.36851210214874963, "updates_nfe_3": 11, "loss_nfe_4": 0.3628325363000234, "updates_nfe_4": 9, "loss_nfe_2": 0.3695180184311337, "updates_nfe_2": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48780325725656, "grad_norm": 0.07644847326805493, "grad_rms": 0.00022261048699289692, "grad_to_weight": 0.004636668213178046, "update_to_weight": 7.494512133288542e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248123764881477, "grad_norm": 0.05348396243452536, "grad_rms": 0.00015574007452199342, "grad_to_weight": 0.0032917008270287205, "update_to_weight": 6.288146443454177e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060513330287316, "grad_norm": 0.0032158821348370237, "grad_rms": 5.7573525747861946e-05, "grad_to_weight": 0.0006354853598724616, "update_to_weight": 3.6927974383108135e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.845417295027765, "grad_norm": 0.02839163009438219, "grad_rms": 9.085333259436315e-05, "grad_to_weight": 0.002210253621372977, "update_to_weight": 8.275051501970339e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.054479579982779, "grad_norm": 0.025855798571238296, "grad_rms": 0.0005882397856508792, "grad_to_weight": 0.005115422500396448, "update_to_weight": 2.9442609106771572e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638130681603776, "grad_norm": 0.08640463618863319, "grad_rms": 0.00020820454707915474, "grad_to_weight": 0.005193169704104314, "update_to_weight": 8.709738659823282e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56002740814966, "grad_norm": 0.0329789489749251, "grad_rms": 4.549941670830103e-05, "grad_to_weight": 0.0013427895835320535, "update_to_weight": 0.00011847228359835999, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826151023344813, "grad_norm": 0.020243178339913636, "grad_rms": 2.4574746912217844e-05, "grad_to_weight": 0.0006566884825998343, "update_to_weight": 0.0001121097312991838, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.603135018737824, "grad_norm": 0.013453094960575427, "grad_rms": 1.2076873448357957e-05, "grad_to_weight": 0.00038878254682098856, "update_to_weight": 0.00013353094059468377, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.03988394865307, "grad_norm": 0.010868272682454033, "grad_rms": 8.742023541966259e-06, "grad_to_weight": 0.000293420808162366, "update_to_weight": 0.00013998674505535873, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66294640331866, "grad_norm": 0.00432684885544932, "grad_rms": 7.81728047743442e-06, "grad_to_weight": 0.00017543925144593013, "update_to_weight": 9.311654856353147e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.29750339856041, "grad_norm": 0.010291109755858046, "grad_rms": 8.277775722715465e-06, "grad_to_weight": 0.0002759195339668568, "update_to_weight": 0.00014046050500893064, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.413513399279694, "grad_norm": 0.022463587158899494, "grad_rms": 9.973054378143535e-06, "grad_to_weight": 0.00035423975040554457, "update_to_weight": 0.00013714451273692208, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75716015145671, "grad_norm": 0.03191889393921806, "grad_rms": 2.3816953252808995e-05, "grad_to_weight": 0.00062885499984581, "update_to_weight": 9.771507802810117e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.071120550197215, "grad_norm": 0.04386925606041346, "grad_rms": 3.767318571893366e-05, "grad_to_weight": 0.0010947848589723942, "update_to_weight": 0.00010006990001925135, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198410054289297, "grad_norm": 0.0916594023788342, "grad_rms": 0.0001428840193702237, "grad_to_weight": 0.00325051668524452, "update_to_weight": 7.144129594384146e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40723013570388, "grad_norm": 0.0038354855588202735, "grad_rms": 0.00032649821305766265, "grad_to_weight": 0.0004562127474697889, "update_to_weight": 5.233925194409761e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3097166339507387, "grad_norm": 0.05035599675714373, "grad_rms": 0.0011657230319627283, "grad_to_weight": 0.038448008868334896, "update_to_weight": 7.748510004912969e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43250}
{"step": 43300, "lr": 1.4347642562002103e-05, "window_steps": 50, "elapsed_seconds": 1339.479024145, "steps_per_second": 1.7170854925989663, "eta_seconds": 3901.9606355528263, "images_seen": 44339200, "images_per_second": 1758.2955444213414, "effective_batch_size": 1024, "loss": 0.36836692690849304, "edge_loss": 0.1928967982530594, "diag_loss": 0.17547008395195007, "jvp_rms": 2.232734441757202, "grad_norm": 0.21698571741580963, "grad_norm_postclip": 0.21698571741580963, "clip_fraction": 0.0, "ema_updates": 43300, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3365227257212003, "updates_nfe_5": 6, "loss_nfe_3": 0.34313300022712123, "updates_nfe_3": 13, "loss_nfe_1": 0.40103008784353733, "updates_nfe_1": 16, "loss_nfe_2": 0.37372535738078033, "updates_nfe_2": 11, "loss_nfe_4": 0.35275528207421303, "updates_nfe_4": 4, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.485756940000226, "grad_norm": 0.050621045938359296, "grad_rms": 0.00014740354132273804, "grad_to_weight": 0.003070592762139717, "update_to_weight": 9.544618792872971e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248520064820518, "grad_norm": 0.014501163612895709, "grad_rms": 4.2225972028395e-05, "grad_to_weight": 0.0008924605782585708, "update_to_weight": 8.864366029329194e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0609090683536815, "grad_norm": 0.0020457183268116803, "grad_rms": 3.662423304812265e-05, "grad_to_weight": 0.0004042195382651193, "update_to_weight": 4.715784545817677e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.844548332853126, "grad_norm": 0.012528158115303709, "grad_rms": 4.009015728440603e-05, "grad_to_weight": 0.0009753677428469656, "update_to_weight": 9.491804487519301e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.054391843471243, "grad_norm": 0.027709539763152862, "grad_rms": 0.000630413858069475, "grad_to_weight": 0.005482269800459826, "update_to_weight": 3.25714861503794e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.637734490740467, "grad_norm": 0.06778816261059105, "grad_rms": 0.00016334544436775162, "grad_to_weight": 0.0040743625671101884, "update_to_weight": 9.466083763144185e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.5599615517977, "grad_norm": 0.03511801766686235, "grad_rms": 4.845058346186051e-05, "grad_to_weight": 0.0014298889512834698, "update_to_weight": 0.00011318429586494719, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826622254934737, "grad_norm": 0.015695920382441716, "grad_rms": 1.905448168642041e-05, "grad_to_weight": 0.0005091677009773299, "update_to_weight": 9.358513207137015e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60522336837316, "grad_norm": 0.010580569176587925, "grad_rms": 9.498200624593212e-06, "grad_to_weight": 0.000305750639548186, "update_to_weight": 0.00010370843251378514, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.04206395374217, "grad_norm": 0.008519641485822686, "grad_rms": 6.852874289620523e-06, "grad_to_weight": 0.0002299991030862088, "update_to_weight": 0.00010599089139648064, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663188929730758, "grad_norm": 0.003387213147667425, "grad_rms": 6.119648755195558e-06, "grad_to_weight": 0.00013733881524072655, "update_to_weight": 7.212618686897099e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.2984901593089, "grad_norm": 0.008588085403965282, "grad_rms": 6.907927963850924e-06, "grad_to_weight": 0.00023025289676026958, "update_to_weight": 0.00011110980973829624, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.415725173010124, "grad_norm": 0.02033894773140294, "grad_rms": 9.02978808703489e-06, "grad_to_weight": 0.0003207240424344977, "update_to_weight": 0.0001223900707478376, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7574295092908, "grad_norm": 0.030247609018442683, "grad_rms": 2.2569888899450425e-05, "grad_to_weight": 0.0005959247603920932, "update_to_weight": 0.00010066442204334478, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07142061369924, "grad_norm": 0.03139630172284807, "grad_rms": 2.6961904803301654e-05, "grad_to_weight": 0.0007835085764869189, "update_to_weight": 0.00011297543196824874, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.197959836925147, "grad_norm": 0.05541147470166239, "grad_rms": 8.637863676965521e-05, "grad_to_weight": 0.0019650880780779476, "update_to_weight": 9.096404726923974e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407575424813373, "grad_norm": 0.005723889371826342, "grad_rms": 0.00048724982091077646, "grad_to_weight": 0.0006808014299738973, "update_to_weight": 5.378507792050526e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.310224786284001, "grad_norm": 0.08447021965966872, "grad_rms": 0.0019554548993860883, "grad_to_weight": 0.0644700211322053, "update_to_weight": 7.765828652897268e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43300}
{"step": 43350, "lr": 1.428404995185735e-05, "window_steps": 50, "elapsed_seconds": 1366.941767539, "steps_per_second": 1.7191661384602124, "eta_seconds": 3868.154363461426, "images_seen": 44390400, "images_per_second": 1760.4261257832575, "effective_batch_size": 1024, "loss": 0.34791266918182373, "edge_loss": 0.17979809641838074, "diag_loss": 0.1681145578622818, "jvp_rms": 1.6523046493530273, "grad_norm": 0.25235193967819214, "grad_norm_postclip": 0.2389621138572693, "clip_fraction": 0.019999999552965164, "ema_updates": 43350, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3578481939103868, "updates_nfe_2": 9, "loss_nfe_5": 0.33991117278734845, "updates_nfe_5": 15, "loss_nfe_4": 0.3297087318367428, "updates_nfe_4": 9, "loss_nfe_1": 0.4064215049147606, "updates_nfe_1": 8, "loss_nfe_3": 0.3175090965297487, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.487433069424423, "grad_norm": 0.06372759770769282, "grad_rms": 0.00018556853988246905, "grad_to_weight": 0.0038652225267178934, "update_to_weight": 8.700612540330799e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.250799071399054, "grad_norm": 0.022812878386334723, "grad_rms": 6.642887359548853e-05, "grad_to_weight": 0.0014038004092047844, "update_to_weight": 7.698210949706757e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060598322357646, "grad_norm": 0.002951548056222967, "grad_rms": 5.284118661258607e-05, "grad_to_weight": 0.0005832409269044478, "update_to_weight": 4.3513621848347116e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.846494061337076, "grad_norm": 0.01656758375886385, "grad_rms": 5.301633588931769e-05, "grad_to_weight": 0.0012896579938277324, "update_to_weight": 9.521925236207023e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053882186863026, "grad_norm": 0.029149169912289934, "grad_rms": 0.0006631665780449066, "grad_to_weight": 0.005767678951452367, "update_to_weight": 3.232392629539542e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638242567129986, "grad_norm": 0.057732767220153414, "grad_rms": 0.00013911550561310752, "grad_to_weight": 0.0034698837324447077, "update_to_weight": 8.893194910092786e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.560739391019844, "grad_norm": 0.02634191516082594, "grad_rms": 3.634263104347018e-05, "grad_to_weight": 0.0010725212601073951, "update_to_weight": 0.0001050841726640984, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826548996715072, "grad_norm": 0.01529198897126211, "grad_rms": 1.8564118363381278e-05, "grad_to_weight": 0.0004960655496303413, "update_to_weight": 8.895631910460539e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60470834790766, "grad_norm": 0.010120493215778641, "grad_rms": 9.085189405122286e-06, "grad_to_weight": 0.0002924600061364357, "update_to_weight": 0.00010549317048854223, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.042050771935195, "grad_norm": 0.008107088829051841, "grad_rms": 6.521032685791984e-06, "grad_to_weight": 0.000218861770882139, "update_to_weight": 0.00010881764068317712, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.6640625, "grad_norm": 0.0032633762467748186, "grad_rms": 5.895913695322932e-06, "grad_to_weight": 0.00013231300588760746, "update_to_weight": 7.150742881841569e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.2988616196929, "grad_norm": 0.00814086874165193, "grad_rms": 6.548203957605239e-06, "grad_to_weight": 0.000218260514882678, "update_to_weight": 0.00010647286099446272, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.41621795107896, "grad_norm": 0.019469290295835332, "grad_rms": 8.643690317612675e-06, "grad_to_weight": 0.00030700806394437594, "update_to_weight": 0.00011547039040095708, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75723470600029, "grad_norm": 0.02629382028080674, "grad_rms": 1.961968637316374e-05, "grad_to_weight": 0.0005180309848065542, "update_to_weight": 8.668352693493028e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07226443505689, "grad_norm": 0.028128265689625256, "grad_rms": 2.4155444437385665e-05, "grad_to_weight": 0.0007019385124893884, "update_to_weight": 0.00010088462783786989, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198506373897715, "grad_norm": 0.056583478479176975, "grad_rms": 8.82056245756224e-05, "grad_to_weight": 0.0020066126102180415, "update_to_weight": 8.2796326078606e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407215162267942, "grad_norm": 0.005214878475170925, "grad_rms": 0.0004439199359102475, "grad_to_weight": 0.0006202860726790477, "update_to_weight": 3.7075833734453814e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3097678767275092, "grad_norm": 0.0549299621325591, "grad_rms": 0.0012716086688062751, "grad_to_weight": 0.04193870006172629, "update_to_weight": 0.00010852144438702043, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43350}
{"step": 43400, "lr": 1.422090260481905e-05, "window_steps": 50, "elapsed_seconds": 1394.418573578, "steps_per_second": 1.7211474699750553, "eta_seconds": 3834.6510773395, "images_seen": 44441600, "images_per_second": 1762.4550092544566, "effective_batch_size": 1024, "loss": 0.3688916563987732, "edge_loss": 0.19062671065330505, "diag_loss": 0.17826496064662933, "jvp_rms": 2.327265501022339, "grad_norm": 0.1879376918077469, "grad_norm_postclip": 0.1879376918077469, "clip_fraction": 0.0, "ema_updates": 43400, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.40528851250807446, "updates_nfe_4": 6, "loss_nfe_3": 0.39735151330629986, "updates_nfe_3": 9, "loss_nfe_1": 0.397112101316452, "updates_nfe_1": 10, "loss_nfe_2": 0.3663051703146526, "updates_nfe_2": 14, "loss_nfe_5": 0.3033906126564199, "updates_nfe_5": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.488263204059066, "grad_norm": 0.015416457972142321, "grad_rms": 4.489121980043758e-05, "grad_to_weight": 0.0009349958683548375, "update_to_weight": 3.091498138309459e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24917647402429, "grad_norm": 0.012219022583177147, "grad_rms": 3.5580600259742466e-05, "grad_to_weight": 0.0007519779604037355, "update_to_weight": 4.0546448372212616e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06146116547768, "grad_norm": 0.0031651637384968964, "grad_rms": 5.666552079769714e-05, "grad_to_weight": 0.0006253458507367964, "update_to_weight": 3.0517475880119737e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.846771404616062, "grad_norm": 0.007581558465926448, "grad_rms": 2.426101814508773e-05, "grad_to_weight": 0.0005901528272856374, "update_to_weight": 4.4092062567216264e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053967667844945, "grad_norm": 0.02192318234097899, "grad_rms": 0.0004987696684560378, "grad_to_weight": 0.004337816104456249, "update_to_weight": 1.6266219213774364e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638019712428875, "grad_norm": 0.05304845245779991, "grad_rms": 0.00012782796739186202, "grad_to_weight": 0.003188387402749129, "update_to_weight": 6.125238377949138e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56085370397309, "grad_norm": 0.02426971919249308, "grad_rms": 3.348372526281219e-05, "grad_to_weight": 0.0009881464009765704, "update_to_weight": 8.767457010655879e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82679649000099, "grad_norm": 0.015981014837696438, "grad_rms": 1.9400579713435632e-05, "grad_to_weight": 0.00051841309047082, "update_to_weight": 8.400885852815797e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60575072727707, "grad_norm": 0.011577447201989766, "grad_rms": 1.039310025858138e-05, "grad_to_weight": 0.00033455269597327793, "update_to_weight": 0.00010358981751812031, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0434496650344, "grad_norm": 0.00956382073610656, "grad_rms": 7.69277220667878e-06, "grad_to_weight": 0.00025817845860975317, "update_to_weight": 0.00010964561990954463, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66404517737751, "grad_norm": 0.0039158587442162055, "grad_rms": 7.0747481911691095e-06, "grad_to_weight": 0.00015876790348275598, "update_to_weight": 7.380948856229465e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.299732164336646, "grad_norm": 0.009487266754574218, "grad_rms": 7.63119510713955e-06, "grad_to_weight": 0.00025435214153213866, "update_to_weight": 0.00011001465852505439, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.41786949965286, "grad_norm": 0.020830917187248065, "grad_rms": 9.248205479627744e-06, "grad_to_weight": 0.0003284707819987881, "update_to_weight": 0.00012154993358393516, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75785518795342, "grad_norm": 0.027492764290485163, "grad_rms": 2.0514303632947954e-05, "grad_to_weight": 0.0005416455086346938, "update_to_weight": 8.47828520723929e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.072570582456244, "grad_norm": 0.02757358707700688, "grad_rms": 2.367910833634225e-05, "grad_to_weight": 0.0006880912972695239, "update_to_weight": 8.534962807344603e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198378669177284, "grad_norm": 0.04290512098826003, "grad_rms": 6.688300359008198e-05, "grad_to_weight": 0.0015215456708210745, "update_to_weight": 6.196985078866826e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407642575318983, "grad_norm": 0.002457514165949749, "grad_rms": 0.0002091974983196863, "grad_to_weight": 0.00029229527111010803, "update_to_weight": 3.4293680431293207e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3102048606703203, "grad_norm": 0.034237870014243096, "grad_rms": 0.0007925942531419987, "grad_to_weight": 0.026131692105558588, "update_to_weight": 0.00010473430923963108, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43400}
{"step": 43450, "lr": 1.4158201211456465e-05, "window_steps": 50, "elapsed_seconds": 1421.903750502, "steps_per_second": 1.723042082936368, "eta_seconds": 3801.4161493012657, "images_seen": 44492800, "images_per_second": 1764.3950929268408, "effective_batch_size": 1024, "loss": 0.3653281331062317, "edge_loss": 0.18819062411785126, "diag_loss": 0.17713750898838043, "jvp_rms": 1.9311717748641968, "grad_norm": 0.17218324542045593, "grad_norm_postclip": 0.17218324542045593, "clip_fraction": 0.0, "ema_updates": 43450, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3606225624680519, "updates_nfe_3": 12, "loss_nfe_5": 0.36766482483256946, "updates_nfe_5": 11, "loss_nfe_2": 0.33854124943415326, "updates_nfe_2": 9, "loss_nfe_4": 0.36834804548157585, "updates_nfe_4": 9, "loss_nfe_1": 0.3925134374035729, "updates_nfe_1": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48475358718104, "grad_norm": 0.3580670604804913, "grad_rms": 0.0010426563056424695, "grad_to_weight": 0.02172110481280917, "update_to_weight": 4.428239856000614e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251626270786332, "grad_norm": 0.2101432275105348, "grad_rms": 0.0006119165526086053, "grad_to_weight": 0.01293059685283835, "update_to_weight": 6.197925403054952e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060985197270606, "grad_norm": 0.009649015642906617, "grad_rms": 0.00017274509054311475, "grad_to_weight": 0.0019065488767108707, "update_to_weight": 3.0032043442315923e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84619117425655, "grad_norm": 0.11809442769759852, "grad_rms": 0.00037790265234801984, "grad_to_weight": 0.009192952688907265, "update_to_weight": 7.679488283665222e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053157144325134, "grad_norm": 0.05164854193148761, "grad_rms": 0.0011750450155794071, "grad_to_weight": 0.01022104408320067, "update_to_weight": 3.4977075165437135e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63764277854776, "grad_norm": 0.2621023299383284, "grad_rms": 0.0006315737129436559, "grad_to_weight": 0.015753573593747176, "update_to_weight": 8.679791652186313e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56090092003765, "grad_norm": 0.035301689143028184, "grad_rms": 4.8703985868282014e-05, "grad_to_weight": 0.001437312469031941, "update_to_weight": 0.00010459667935455974, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82657572608566, "grad_norm": 0.013398069971503113, "grad_rms": 1.6264944819098967e-05, "grad_to_weight": 0.00043462725443635874, "update_to_weight": 8.532981532920454e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60546147455901, "grad_norm": 0.003350301579929526, "grad_rms": 3.0075732248389604e-06, "grad_to_weight": 9.681424368210134e-05, "update_to_weight": 9.317944117412738e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.04348261829856, "grad_norm": 0.002575457216801705, "grad_rms": 2.071599441644082e-06, "grad_to_weight": 6.952524532694701e-05, "update_to_weight": 9.787915676630589e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663323803081063, "grad_norm": 0.000985011463694361, "grad_rms": 1.7796117087587556e-06, "grad_to_weight": 3.9938309676301966e-05, "update_to_weight": 6.527350695436154e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30058959822217, "grad_norm": 0.002617195881472741, "grad_rms": 2.105172429719133e-06, "grad_to_weight": 7.016500033011495e-05, "update_to_weight": 9.6491077237948e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.41859709132054, "grad_norm": 0.008254378624674813, "grad_rms": 3.664658110895434e-06, "grad_to_weight": 0.00013015706753633795, "update_to_weight": 0.00010521537701472885, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.758249598895254, "grad_norm": 0.03787894271649382, "grad_rms": 2.826416885442566e-05, "grad_to_weight": 0.0007462618001176749, "update_to_weight": 8.244763498899467e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07294831333773, "grad_norm": 0.12770273738406926, "grad_rms": 0.00010966607082784686, "grad_to_weight": 0.003186756721405626, "update_to_weight": 0.0001031984989797294, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198132998300505, "grad_norm": 0.31852915396646786, "grad_rms": 0.0004965418126688056, "grad_to_weight": 0.011296107936850484, "update_to_weight": 8.929270856992578e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407573156205846, "grad_norm": 0.008686226921611978, "grad_rms": 0.0007394207394674657, "grad_to_weight": 0.0010331431865330189, "update_to_weight": 5.4957765999302746e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3102112751132773, "grad_norm": 0.16387632001269542, "grad_rms": 0.0037936772764803352, "grad_to_weight": 0.12507625535318884, "update_to_weight": 0.00013704961145691598, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43450}
{"step": 43500, "lr": 1.4095946457462003e-05, "window_steps": 50, "elapsed_seconds": 1449.387471143, "steps_per_second": 1.7248665727933175, "eta_seconds": 3768.4074249718005, "images_seen": 44544000, "images_per_second": 1766.263370540357, "effective_batch_size": 1024, "loss": 0.38009780645370483, "edge_loss": 0.20053750276565552, "diag_loss": 0.17956028878688812, "jvp_rms": 2.3353514671325684, "grad_norm": 0.2862689197063446, "grad_norm_postclip": 0.2862689197063446, "clip_fraction": 0.0, "ema_updates": 43500, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3360046688467264, "updates_nfe_5": 8, "loss_nfe_3": 0.35760261283980477, "updates_nfe_3": 9, "loss_nfe_2": 0.36029742161432904, "updates_nfe_2": 9, "loss_nfe_1": 0.40586833918795867, "updates_nfe_1": 17, "loss_nfe_4": 0.42228446687970844, "updates_nfe_4": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.488241919077705, "grad_norm": 0.1988251950061206, "grad_rms": 0.0005789595474533186, "grad_to_weight": 0.01205860491263596, "update_to_weight": 0.0001349498569857055, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.250246955755685, "grad_norm": 0.10317377107603316, "grad_rms": 0.0003004319437956222, "grad_to_weight": 0.006349058654735704, "update_to_weight": 0.00010992677240690789, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0613428371033375, "grad_norm": 0.0049596045317303284, "grad_rms": 8.879116436313662e-05, "grad_to_weight": 0.0009798989500123972, "update_to_weight": 5.09305828472833e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.847733448384668, "grad_norm": 0.0698474265257411, "grad_rms": 0.00022351205097797985, "grad_to_weight": 0.005436556323833363, "update_to_weight": 0.0001252022374765133, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053248110600266, "grad_norm": 0.03986807108936577, "grad_rms": 0.000907030023741377, "grad_to_weight": 0.007889593033386584, "update_to_weight": 4.7054525292585105e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638248986768502, "grad_norm": 0.15071286320801955, "grad_rms": 0.0003631645801357629, "grad_to_weight": 0.009058216602472611, "update_to_weight": 0.0001391480019209198, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56180670538738, "grad_norm": 0.02345685358480465, "grad_rms": 3.236225498672234e-05, "grad_to_weight": 0.0009550133614421625, "update_to_weight": 0.00013289808066233597, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826659873952853, "grad_norm": 0.009132056600828348, "grad_rms": 1.1086103969697205e-05, "grad_to_weight": 0.00029623892559779165, "update_to_weight": 9.426514116270954e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.604486110681876, "grad_norm": 0.0035360156463704003, "grad_rms": 3.1742891578312635e-06, "grad_to_weight": 0.0001021837352261355, "update_to_weight": 0.00010756964930219038, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0430707003896, "grad_norm": 0.002926847762959261, "grad_rms": 2.3542445791637745e-06, "grad_to_weight": 7.901201783815612e-05, "update_to_weight": 0.00011152255624292053, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663405469154853, "grad_norm": 0.0010608019331986465, "grad_rms": 1.9165416957827456e-06, "grad_to_weight": 4.301117031568622e-05, "update_to_weight": 7.442865204485543e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30084158840287, "grad_norm": 0.003138773558116522, "grad_rms": 2.5247095964250422e-06, "grad_to_weight": 8.414752655587245e-05, "update_to_weight": 0.00011202102475831564, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.41929195281886, "grad_norm": 0.008735296083693141, "grad_rms": 3.878168799827779e-06, "grad_to_weight": 0.0001377387828642397, "update_to_weight": 0.0001252150080504064, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75863438710453, "grad_norm": 0.029677668075889146, "grad_rms": 2.2144615492060168e-05, "grad_to_weight": 0.0005846821616506865, "update_to_weight": 0.00011910724553792011, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07371137978114, "grad_norm": 0.08758786028690177, "grad_rms": 7.521699758866083e-05, "grad_to_weight": 0.002185668790615024, "update_to_weight": 0.00016058224673269516, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19865355828182, "grad_norm": 0.22871489680224244, "grad_rms": 0.0003565341132149543, "grad_to_weight": 0.008110844595098404, "update_to_weight": 0.0001446886828262719, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407607638974211, "grad_norm": 0.01775974991778933, "grad_rms": 0.0015118103102160316, "grad_to_weight": 0.0021123428542814526, "update_to_weight": 1.2078601559432947e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3102639087534886, "grad_norm": 0.233435504156668, "grad_rms": 0.005403947120452039, "grad_to_weight": 0.17815914992174775, "update_to_weight": 0.00022534358921676343, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43500}
{"step": 43550, "lr": 1.4034139023643683e-05, "window_steps": 50, "elapsed_seconds": 1478.924319684, "steps_per_second": 1.7242261595541652, "eta_seconds": 3740.808573318353, "images_seen": 44595200, "images_per_second": 1765.6075873834652, "effective_batch_size": 1024, "loss": 0.3479771614074707, "edge_loss": 0.180153027176857, "diag_loss": 0.1678241491317749, "jvp_rms": 1.6666796207427979, "grad_norm": 0.27666425704956055, "grad_norm_postclip": 0.27666425704956055, "clip_fraction": 0.0, "ema_updates": 43550, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3039579279720783, "updates_nfe_3": 12, "loss_nfe_5": 0.34471278163519775, "updates_nfe_5": 11, "loss_nfe_2": 0.37113840050167507, "updates_nfe_2": 9, "loss_nfe_1": 0.4074842259287834, "updates_nfe_1": 8, "loss_nfe_4": 0.3359405636787415, "updates_nfe_4": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.484549947235156, "grad_norm": 0.2154759382469374, "grad_rms": 0.0006274448853963692, "grad_to_weight": 0.013071387386167478, "update_to_weight": 8.642390642937197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.250818789467964, "grad_norm": 0.15016245563199673, "grad_rms": 0.00043725840356652774, "grad_to_weight": 0.009240300908980408, "update_to_weight": 8.837625960264256e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062216668383155, "grad_norm": 0.006839670286097668, "grad_rms": 0.0001224497406350021, "grad_to_weight": 0.0013511216003091826, "update_to_weight": 3.787372996891959e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84620661573145, "grad_norm": 0.06909239670995211, "grad_rms": 0.00022109595247484702, "grad_to_weight": 0.0053784279497218765, "update_to_weight": 9.346822894916415e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053911246785094, "grad_norm": 0.02740082043045482, "grad_rms": 0.0006233902500539565, "grad_to_weight": 0.005421705901124618, "update_to_weight": 2.6597711978143264e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638149940641405, "grad_norm": 0.14009850010763328, "grad_rms": 0.000337587727326326, "grad_to_weight": 0.008420317199174876, "update_to_weight": 7.469122025250674e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.561287342838934, "grad_norm": 0.02096691016294081, "grad_rms": 2.892700380823226e-05, "grad_to_weight": 0.0008536568083860596, "update_to_weight": 8.961597996762453e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82681826931503, "grad_norm": 0.0081166131356613, "grad_rms": 9.853379259124301e-06, "grad_to_weight": 0.00026329714162361557, "update_to_weight": 7.484393003252113e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.606819530373095, "grad_norm": 0.0035269172649394277, "grad_rms": 3.1661215204624086e-06, "grad_to_weight": 0.0001019139381428561, "update_to_weight": 8.76960053710636e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.044431659727046, "grad_norm": 0.002974165272382628, "grad_rms": 2.392304976930026e-06, "grad_to_weight": 8.02864327816372e-05, "update_to_weight": 9.387251864608058e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663941241387064, "grad_norm": 0.0010883760347710576, "grad_rms": 1.9663595870717674e-06, "grad_to_weight": 4.4128228498400724e-05, "update_to_weight": 6.208108106237643e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.302508935077846, "grad_norm": 0.003086761801131106, "grad_rms": 2.4828733251692068e-06, "grad_to_weight": 8.274944204164331e-05, "update_to_weight": 9.652627521843191e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.42142268562532, "grad_norm": 0.00842033483389161, "grad_rms": 3.738336917721884e-06, "grad_to_weight": 0.0001327679903308146, "update_to_weight": 0.00010701446776775007, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75912739273475, "grad_norm": 0.02506958559668386, "grad_rms": 1.8706197945345855e-05, "grad_to_weight": 0.0004938931554657127, "update_to_weight": 8.359878304744631e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.074089099909685, "grad_norm": 0.06165912968672456, "grad_rms": 5.295042707715105e-05, "grad_to_weight": 0.0015386283524249469, "update_to_weight": 8.939322013091212e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198390573879056, "grad_norm": 0.17893778885282757, "grad_rms": 0.0002789386557730424, "grad_to_weight": 0.006345673820781196, "update_to_weight": 7.118089784545898e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40815072392208, "grad_norm": 0.003709975363168826, "grad_rms": 0.00031581407568515235, "grad_to_weight": 0.0004412355921039275, "update_to_weight": 5.502127035884261e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.311010832889474, "grad_norm": 0.06219602963643037, "grad_rms": 0.001439815491956033, "grad_to_weight": 0.04744127819245401, "update_to_weight": 8.822317524671014e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43550}
{"step": 43600, "lr": 1.3972779585917677e-05, "window_steps": 50, "elapsed_seconds": 1506.418711083, "steps_per_second": 1.7259477599895174, "eta_seconds": 3708.1075965120003, "images_seen": 44646400, "images_per_second": 1767.3705062292659, "effective_batch_size": 1024, "loss": 0.3570559620857239, "edge_loss": 0.18376274406909943, "diag_loss": 0.1732931286096573, "jvp_rms": 1.8800389766693115, "grad_norm": 0.2236190140247345, "grad_norm_postclip": 0.2236190140247345, "clip_fraction": 0.0, "ema_updates": 43600, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3972770541906357, "updates_nfe_4": 10, "loss_nfe_3": 0.3342583099236855, "updates_nfe_3": 13, "loss_nfe_2": 0.350301899693229, "updates_nfe_2": 11, "loss_nfe_5": 0.3229649648070335, "updates_nfe_5": 10, "loss_nfe_1": 0.40861599644025165, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.485913361275095, "grad_norm": 0.002948669534575325, "grad_rms": 8.58623767110886e-06, "grad_to_weight": 0.00017885994363537415, "update_to_weight": 7.843600764536082e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.249961500855836, "grad_norm": 0.005097483874049717, "grad_rms": 1.484337514055783e-05, "grad_to_weight": 0.00031369205851849237, "update_to_weight": 6.173922308944408e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061778264344231, "grad_norm": 0.0018757519246104346, "grad_rms": 3.358134633054063e-05, "grad_to_weight": 0.00037057172927219167, "update_to_weight": 3.75072084784093e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84658789549352, "grad_norm": 0.005206992250195041, "grad_rms": 1.6662396528505334e-05, "grad_to_weight": 0.0004053210309658654, "update_to_weight": 8.73027724200241e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053412864851001, "grad_norm": 0.01255739588423453, "grad_rms": 0.0002856906485763044, "grad_to_weight": 0.0024849336913628137, "update_to_weight": 2.645346550297259e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638575467975485, "grad_norm": 0.04068429984973885, "grad_rms": 9.8034742082066e-05, "grad_to_weight": 0.002445179271990233, "update_to_weight": 8.755457515816092e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56178309823706, "grad_norm": 0.02247629081462761, "grad_rms": 3.100942127080092e-05, "grad_to_weight": 0.000915091983539292, "update_to_weight": 0.00011115750711910283, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826499497819444, "grad_norm": 0.014046970955724629, "grad_rms": 1.7052695496910804e-05, "grad_to_weight": 0.00045567843201652726, "update_to_weight": 0.00010179990951402837, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60648090326164, "grad_norm": 0.010824282641688603, "grad_rms": 9.716982747540169e-06, "grad_to_weight": 0.00031278195179528987, "update_to_weight": 0.00012026376651167649, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.044046114579096, "grad_norm": 0.009216815021950321, "grad_rms": 7.41365406058662e-06, "grad_to_weight": 0.00024880692010376647, "update_to_weight": 0.00012177212299673163, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663467337212545, "grad_norm": 0.0036385337362220336, "grad_rms": 6.573707493117719e-06, "grad_to_weight": 0.00014752725910246079, "update_to_weight": 8.379290055418647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30295398459165, "grad_norm": 0.008944260815079938, "grad_rms": 7.194421864032565e-06, "grad_to_weight": 0.00023977352621388782, "update_to_weight": 0.00011958845478964417, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.421651730328854, "grad_norm": 0.01930065298187407, "grad_rms": 8.568821193174843e-06, "grad_to_weight": 0.00030432277393123, "update_to_weight": 0.00012805115701321776, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.75958191563441, "grad_norm": 0.025751641052446973, "grad_rms": 1.9215127952026763e-05, "grad_to_weight": 0.0005073257123206374, "update_to_weight": 9.046823707665158e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07465871949972, "grad_norm": 0.02483182158898461, "grad_rms": 2.1324588344351026e-05, "grad_to_weight": 0.0006196390033610398, "update_to_weight": 9.782828320829468e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198809399733527, "grad_norm": 0.031316232528766155, "grad_rms": 4.88175687285074e-05, "grad_to_weight": 0.0011105515869426065, "update_to_weight": 7.910275293345349e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.407818162282055, "grad_norm": 0.002423692681956042, "grad_rms": 0.00020631842240672674, "grad_to_weight": 0.0002882665437305559, "update_to_weight": 5.449203182814781e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3106331036507715, "grad_norm": 0.030136131576336024, "grad_rms": 0.0006976404983545557, "grad_to_weight": 0.022993568140764767, "update_to_weight": 0.00010028722478804814, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43600}
{"step": 43650, "lr": 1.3911868815300991e-05, "window_steps": 50, "elapsed_seconds": 1533.900961641, "steps_per_second": 1.7276213173274064, "eta_seconds": 3675.574002422774, "images_seen": 44697600, "images_per_second": 1769.0842289432642, "effective_batch_size": 1024, "loss": 0.3541712760925293, "edge_loss": 0.18345743417739868, "diag_loss": 0.17071378231048584, "jvp_rms": 1.7786718606948853, "grad_norm": 0.23267585039138794, "grad_norm_postclip": 0.21758128702640533, "clip_fraction": 0.019999999552965164, "ema_updates": 43650, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.34911543428897857, "updates_nfe_3": 10, "loss_nfe_4": 0.33567335388877173, "updates_nfe_4": 11, "loss_nfe_2": 0.3536688178777695, "updates_nfe_2": 10, "loss_nfe_5": 0.3325183358457353, "updates_nfe_5": 9, "loss_nfe_1": 0.3995646834373474, "updates_nfe_1": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.489219149023743, "grad_norm": 0.007040506437837286, "grad_rms": 2.0501267060077424e-05, "grad_to_weight": 0.0004269763397652535, "update_to_weight": 3.40566303524351e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.249826283626753, "grad_norm": 0.006620407654001998, "grad_rms": 1.9277980435022284e-05, "grad_to_weight": 0.0004074140571381178, "update_to_weight": 4.42979322697547e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061280280220724, "grad_norm": 0.0007222431138798173, "grad_rms": 1.2930225913047328e-05, "grad_to_weight": 0.00014269968740958957, "update_to_weight": 2.4712777087884373e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84773701137584, "grad_norm": 0.005691101324105287, "grad_rms": 1.82115475479327e-05, "grad_to_weight": 0.0004429652723328774, "update_to_weight": 5.722009081926174e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053396823722417, "grad_norm": 0.014849106863546326, "grad_rms": 0.0003378288786731199, "grad_to_weight": 0.00293844069277113, "update_to_weight": 2.4090597685750984e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63869285291112, "grad_norm": 0.05038857862472112, "grad_rms": 0.00012141861424679642, "grad_to_weight": 0.00302839766742284, "update_to_weight": 6.486383716922405e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56180422042526, "grad_norm": 0.015450075398716335, "grad_rms": 2.131570109391157e-05, "grad_to_weight": 0.000629028521686052, "update_to_weight": 8.146293083043331e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82661829503542, "grad_norm": 0.009534471622783192, "grad_rms": 1.157462643154383e-05, "grad_to_weight": 0.0003092934661703942, "update_to_weight": 7.511855913037082e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60742093806435, "grad_norm": 0.006112052192933697, "grad_rms": 5.486802929744706e-06, "grad_to_weight": 0.00017661102813388538, "update_to_weight": 9.454659792752966e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.04524392761168, "grad_norm": 0.005067665591466849, "grad_rms": 4.076236693521417e-06, "grad_to_weight": 0.00013679665873895527, "update_to_weight": 0.0001009686409157105, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663555189587807, "grad_norm": 0.0017362149060112615, "grad_rms": 3.1368044835444203e-06, "grad_to_weight": 7.039597059973892e-05, "update_to_weight": 6.529156344770707e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30349392724075, "grad_norm": 0.005264734684117493, "grad_rms": 4.234751546587916e-06, "grad_to_weight": 0.0001411324819703534, "update_to_weight": 9.97463417511343e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.42337050158443, "grad_norm": 0.013289016843814902, "grad_rms": 5.8998630395914114e-06, "grad_to_weight": 0.00020952870745781193, "update_to_weight": 0.00010823781657759543, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76001960050597, "grad_norm": 0.021040458576528843, "grad_rms": 1.5699780176879392e-05, "grad_to_weight": 0.00041450848014091613, "update_to_weight": 8.084252299708707e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.075263360312896, "grad_norm": 0.024813463036677873, "grad_rms": 2.130882274418607e-05, "grad_to_weight": 0.0006191715526254282, "update_to_weight": 9.135436523701631e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198791001829207, "grad_norm": 0.05482201339409243, "grad_rms": 8.545975012297255e-05, "grad_to_weight": 0.0019441263772810762, "update_to_weight": 7.128920715120231e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408038207950037, "grad_norm": 0.006636997450046955, "grad_rms": 0.0005649787423981632, "grad_to_weight": 0.0007893633789355865, "update_to_weight": 6.481049381790576e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3109695503399916, "grad_norm": 0.08576872899267113, "grad_rms": 0.0019855149187319415, "grad_to_weight": 0.06542389102052565, "update_to_weight": 0.00020789364370466039, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43650}
{"step": 43700, "lr": 1.3851407377904024e-05, "window_steps": 50, "elapsed_seconds": 1561.4364153640001, "steps_per_second": 1.729177040725401, "eta_seconds": 3643.351635849334, "images_seen": 44748800, "images_per_second": 1770.6772897028106, "effective_batch_size": 1024, "loss": 0.3683198392391205, "edge_loss": 0.1879325658082962, "diag_loss": 0.18038736283779144, "jvp_rms": 1.9049218893051147, "grad_norm": 0.2270020842552185, "grad_norm_postclip": 0.2270020842552185, "clip_fraction": 0.0, "ema_updates": 43700, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3759911259015401, "updates_nfe_4": 15, "loss_nfe_2": 0.35833199322223663, "updates_nfe_2": 10, "loss_nfe_3": 0.3503626320097182, "updates_nfe_3": 9, "loss_nfe_1": 0.39448037147521975, "updates_nfe_1": 5, "loss_nfe_5": 0.36974034119736066, "updates_nfe_5": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.48899613100213, "grad_norm": 0.0405717041066934, "grad_rms": 0.00011814083948615296, "grad_to_weight": 0.002460532089665038, "update_to_weight": 4.499075512060142e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248461841264582, "grad_norm": 0.04481987412700278, "grad_rms": 0.00013051109564200567, "grad_to_weight": 0.0027584071997004823, "update_to_weight": 6.663340882346469e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06086761153118, "grad_norm": 0.0030605998691944263, "grad_rms": 5.479352724533179e-05, "grad_to_weight": 0.0006047579395716366, "update_to_weight": 3.2003315440087725e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84705111736189, "grad_norm": 0.016661783694553303, "grad_rms": 5.3317776069367615e-05, "grad_to_weight": 0.0012969344904400723, "update_to_weight": 6.926209107851021e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053433435164906, "grad_norm": 0.018074393923015495, "grad_rms": 0.0004112067000270954, "grad_to_weight": 0.0035766561793893863, "update_to_weight": 1.8710935216037894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638764383950193, "grad_norm": 0.033077808164517115, "grad_rms": 7.970579324272942e-05, "grad_to_weight": 0.001987996668576188, "update_to_weight": 5.52570460719136e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.562065140075507, "grad_norm": 0.019452670842014894, "grad_rms": 2.6837883081210683e-05, "grad_to_weight": 0.0007919802643254082, "update_to_weight": 7.383089426674791e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826916276037792, "grad_norm": 0.012262557225613015, "grad_rms": 1.4886458798905869e-05, "grad_to_weight": 0.0003977873464802212, "update_to_weight": 6.437089951492543e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60820222335758, "grad_norm": 0.008735655072474424, "grad_rms": 7.842017105205014e-06, "grad_to_weight": 0.00025241574283736137, "update_to_weight": 7.953979322952732e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.04649607039443, "grad_norm": 0.007021396180296681, "grad_rms": 5.647742976187989e-06, "grad_to_weight": 0.00018952929224277715, "update_to_weight": 8.651072229769077e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.663937529379307, "grad_norm": 0.0030723116546641973, "grad_rms": 5.550718946041534e-06, "grad_to_weight": 0.00012456695736455327, "update_to_weight": 5.8039069281491044e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30457869748531, "grad_norm": 0.00737213119606005, "grad_rms": 5.929860830089953e-06, "grad_to_weight": 0.00019762000948577934, "update_to_weight": 9.06435854805144e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.42511039727286, "grad_norm": 0.018275152807379225, "grad_rms": 8.113534647322305e-06, "grad_to_weight": 0.0002881375009504913, "update_to_weight": 0.00010352928298854097, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76064245323955, "grad_norm": 0.026266377731273268, "grad_rms": 1.9599209538326628e-05, "grad_to_weight": 0.0005174555809743685, "update_to_weight": 7.54402225315259e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0755512087794, "grad_norm": 0.032477238363125424, "grad_rms": 2.789017052869917e-05, "grad_to_weight": 0.000810400291038557, "update_to_weight": 8.217801707328465e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19882455093984, "grad_norm": 0.052138232811427536, "grad_rms": 8.127611650976157e-05, "grad_to_weight": 0.0018489505730014486, "update_to_weight": 6.081226685001765e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408259608913067, "grad_norm": 0.006273879116999471, "grad_rms": 0.0005340680571536731, "grad_to_weight": 0.0007461566850705854, "update_to_weight": 5.492365487875962e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.311242045695686, "grad_norm": 0.07119605450833223, "grad_rms": 0.0016481627982761098, "grad_to_weight": 0.05429665311758579, "update_to_weight": 0.00011731751706667596, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43700}
{"step": 43750, "lr": 1.3791395934923398e-05, "window_steps": 50, "elapsed_seconds": 1588.944717716, "steps_per_second": 1.7307084188258846, "eta_seconds": 3611.237994809091, "images_seen": 44800000, "images_per_second": 1772.2454208777058, "effective_batch_size": 1024, "loss": 0.3446356952190399, "edge_loss": 0.1772148162126541, "diag_loss": 0.16742093861103058, "jvp_rms": 1.5972265005111694, "grad_norm": 0.1716325283050537, "grad_norm_postclip": 0.1716325283050537, "clip_fraction": 0.0, "ema_updates": 43750, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.31666784394871106, "updates_nfe_5": 11, "loss_nfe_2": 0.35651834444566205, "updates_nfe_2": 11, "loss_nfe_4": 0.33365317720633286, "updates_nfe_4": 13, "loss_nfe_3": 0.33634358271956444, "updates_nfe_3": 8, "loss_nfe_1": 0.3997856889452253, "updates_nfe_1": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.490617338507317, "grad_norm": 0.015007217339978625, "grad_rms": 4.3699551052471757e-05, "grad_to_weight": 0.0009100458176865946, "update_to_weight": 3.230301976013883e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24928070823925, "grad_norm": 0.006211708592380433, "grad_rms": 1.808788868757683e-05, "grad_to_weight": 0.0003822759114026977, "update_to_weight": 4.8143057283523915e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06134811294965, "grad_norm": 0.0007718385313294961, "grad_rms": 1.3818126316045074e-05, "grad_to_weight": 0.00015249663016750776, "update_to_weight": 2.9355866715872574e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.8477542324859, "grad_norm": 0.007928384607195936, "grad_rms": 2.5370863217752696e-05, "grad_to_weight": 0.0006171027608193809, "update_to_weight": 5.876636805281307e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.053110716922841, "grad_norm": 0.020404105709873253, "grad_rms": 0.000464209478431085, "grad_to_weight": 0.004037929674000613, "update_to_weight": 2.4102064088832874e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638444326013285, "grad_norm": 0.07432151255565543, "grad_rms": 0.00017908850198854976, "grad_to_weight": 0.004466854658969401, "update_to_weight": 8.997401607053732e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56215832499317, "grad_norm": 0.016742340817463727, "grad_rms": 2.3098575461133523e-05, "grad_to_weight": 0.0006816314998029955, "update_to_weight": 0.0001108880743185203, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826769760821822, "grad_norm": 0.006778766683419259, "grad_rms": 8.229264833059478e-06, "grad_to_weight": 0.00021989870284866788, "update_to_weight": 0.00010146256630392723, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.607651974604, "grad_norm": 0.0038107029745572114, "grad_rms": 3.420876586976846e-06, "grad_to_weight": 0.00011011157235843693, "update_to_weight": 0.00011772973870415431, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.046654232788015, "grad_norm": 0.003116055343927824, "grad_rms": 2.5064359324241697e-06, "grad_to_weight": 8.41116534936634e-05, "update_to_weight": 0.0001174448793734146, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.664470815435486, "grad_norm": 0.0011039124188923926, "grad_rms": 1.9944290381524714e-06, "grad_to_weight": 4.475719050098345e-05, "update_to_weight": 8.182415989034059e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30477012425405, "grad_norm": 0.0033284766590919385, "grad_rms": 2.6772995270575735e-06, "grad_to_weight": 8.922388874145342e-05, "update_to_weight": 0.00011227362398383787, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.42586485047918, "grad_norm": 0.008873180180632918, "grad_rms": 3.939384561448327e-06, "grad_to_weight": 0.00013989844997069648, "update_to_weight": 0.00011428269737535279, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.761020008898065, "grad_norm": 0.021385495676706494, "grad_rms": 1.5957236857585957e-05, "grad_to_weight": 0.0004212975955360582, "update_to_weight": 8.488343606042375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07581164132471, "grad_norm": 0.02958084211877285, "grad_rms": 2.5402859745975853e-05, "grad_to_weight": 0.0007381220967779519, "update_to_weight": 9.130591125653481e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19879965966803, "grad_norm": 0.08287538455299266, "grad_rms": 0.00012919098035183102, "grad_to_weight": 0.002938968521824248, "update_to_weight": 6.594171133062064e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408248266792295, "grad_norm": 0.013832994380093404, "grad_rms": 0.0011775426806003577, "grad_to_weight": 0.0016451695931393594, "update_to_weight": 4.653885060504112e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3113364557178315, "grad_norm": 0.1916664262611153, "grad_rms": 0.004437008140740878, "grad_to_weight": 0.14616113616409468, "update_to_weight": 0.0001063114680650775, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43750}
{"step": 43800, "lr": 1.3731835142634623e-05, "window_steps": 50, "elapsed_seconds": 1616.447874497, "steps_per_second": 1.732193189880183, "eta_seconds": 3579.277436386214, "images_seen": 44851200, "images_per_second": 1773.7658264373074, "effective_batch_size": 1024, "loss": 0.37409743666648865, "edge_loss": 0.19447267055511475, "diag_loss": 0.1796247959136963, "jvp_rms": 2.2279295921325684, "grad_norm": 0.28524544835090637, "grad_norm_postclip": 0.28524544835090637, "clip_fraction": 0.0, "ema_updates": 43800, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39772092608305126, "updates_nfe_1": 13, "loss_nfe_3": 0.39243874152501423, "updates_nfe_3": 15, "loss_nfe_2": 0.3619161956012249, "updates_nfe_2": 8, "loss_nfe_4": 0.295057138800621, "updates_nfe_4": 5, "loss_nfe_5": 0.36414490309026504, "updates_nfe_5": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.488869351777883, "grad_norm": 0.1466229165104794, "grad_rms": 0.00042695161137188705, "grad_to_weight": 0.008892235930941503, "update_to_weight": 7.553780652941043e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248504100317877, "grad_norm": 0.05402947599999588, "grad_rms": 0.00015732855674118466, "grad_to_weight": 0.0033251969329865187, "update_to_weight": 8.636800948945135e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06152258962587, "grad_norm": 0.006755796547655747, "grad_rms": 0.00012094815984401326, "grad_to_weight": 0.0013347360261717436, "update_to_weight": 3.987304133332587e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.847311820786006, "grad_norm": 0.03075276350429692, "grad_rms": 9.840896917731129e-05, "grad_to_weight": 0.0023937119245865278, "update_to_weight": 9.108194080507542e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.052916321134625, "grad_norm": 0.062340891492892535, "grad_rms": 0.0014183043911030725, "grad_to_weight": 0.012337606152736361, "update_to_weight": 3.1588415654124614e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63799770192606, "grad_norm": 0.11422616569252354, "grad_rms": 0.00027524457183848434, "grad_to_weight": 0.006865379340646285, "update_to_weight": 8.512989406061748e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.561640212369266, "grad_norm": 0.05390940541209652, "grad_rms": 7.437612712299268e-05, "grad_to_weight": 0.0021948617822741207, "update_to_weight": 0.00011381422692283125, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826565826321477, "grad_norm": 0.046243879341746834, "grad_rms": 5.61389922066663e-05, "grad_to_weight": 0.001500130750933702, "update_to_weight": 0.00011192482302003348, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60960249643063, "grad_norm": 0.034904792659171985, "grad_rms": 3.1334110472075595e-05, "grad_to_weight": 0.0010085291405115616, "update_to_weight": 0.0001335337730009508, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0483972694417, "grad_norm": 0.028743493547975722, "grad_rms": 2.312016864284467e-05, "grad_to_weight": 0.0007758363563997939, "update_to_weight": 0.0001396915751608777, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665177308904298, "grad_norm": 0.009996042060876797, "grad_rms": 1.8059762904750435e-05, "grad_to_weight": 0.0004052694183255742, "update_to_weight": 9.308673701619839e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.305534185417784, "grad_norm": 0.02649041472163156, "grad_rms": 2.130787806850117e-05, "grad_to_weight": 0.0007100934298371821, "update_to_weight": 0.00013634834441558124, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.427647018078794, "grad_norm": 0.03702409884459222, "grad_rms": 1.6437417072659953e-05, "grad_to_weight": 0.0005837217772565209, "update_to_weight": 0.0001281726703767038, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.761344656422615, "grad_norm": 0.040330109596655825, "grad_rms": 3.009315851524604e-05, "grad_to_weight": 0.0007945043589690059, "update_to_weight": 8.89146244094138e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.076166496981095, "grad_norm": 0.047631768579010136, "grad_rms": 4.090428297500984e-05, "grad_to_weight": 0.0011885310582936668, "update_to_weight": 0.00010207979656489778, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.198617844494503, "grad_norm": 0.13834519928094863, "grad_rms": 0.0002156605609551863, "grad_to_weight": 0.0049060985912101765, "update_to_weight": 8.503214668010573e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408469208551123, "grad_norm": 0.010277365772703275, "grad_rms": 0.0008748674733010096, "grad_to_weight": 0.0012222635913623312, "update_to_weight": 6.600446888554594e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3116347776234114, "grad_norm": 0.09122185488595415, "grad_rms": 0.0021117528021890003, "grad_to_weight": 0.06954821299511564, "update_to_weight": 0.0002255148749321358, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43800}
{"step": 43850, "lr": 1.3672725652384976e-05, "window_steps": 50, "elapsed_seconds": 1643.916989019, "steps_per_second": 1.7336641807569155, "eta_seconds": 3547.399818409421, "images_seen": 44902400, "images_per_second": 1775.2721210950815, "effective_batch_size": 1024, "loss": 0.3717174828052521, "edge_loss": 0.19297806918621063, "diag_loss": 0.1787392944097519, "jvp_rms": 2.049570322036743, "grad_norm": 0.2518162727355957, "grad_norm_postclip": 0.2518162727355957, "clip_fraction": 0.0, "ema_updates": 43850, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.37376851836840314, "updates_nfe_3": 9, "loss_nfe_5": 0.3988527297973633, "updates_nfe_5": 10, "loss_nfe_4": 0.3388953777876767, "updates_nfe_4": 11, "loss_nfe_2": 0.3472692251205444, "updates_nfe_2": 10, "loss_nfe_1": 0.40328856408596037, "updates_nfe_1": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.488145673601572, "grad_norm": 0.025937255271509675, "grad_rms": 7.552675390919231e-05, "grad_to_weight": 0.0015730850384853554, "update_to_weight": 7.690818527575797e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25218209507765, "grad_norm": 0.02268651490562114, "grad_rms": 6.606091548230116e-05, "grad_to_weight": 0.0013959057788610599, "update_to_weight": 9.707562947209138e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061395406897768, "grad_norm": 0.0015161340198772367, "grad_rms": 2.714315306678216e-05, "grad_to_weight": 0.00029954862206794193, "update_to_weight": 4.336951967531694e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.848368239050824, "grad_norm": 0.013148251044547307, "grad_rms": 4.207445719789106e-05, "grad_to_weight": 0.0010233401471624257, "update_to_weight": 0.00010898539595204402, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.051929694392526, "grad_norm": 0.05661408649454389, "grad_rms": 0.0012880150660446567, "grad_to_weight": 0.011206428022421541, "update_to_weight": 4.119783280830065e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638076572759715, "grad_norm": 0.09193627446166244, "grad_rms": 0.00022153383462719127, "grad_to_weight": 0.0055256552077770135, "update_to_weight": 0.0001126267322462207, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.562053957861632, "grad_norm": 0.024554091850948888, "grad_rms": 3.387606008516708e-05, "grad_to_weight": 0.0009996758370889338, "update_to_weight": 0.00011916951385200626, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826545036806348, "grad_norm": 0.01153556568885785, "grad_rms": 1.400390800954407e-05, "grad_to_weight": 0.00037420884095459254, "update_to_weight": 8.49613756109176e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.60865546656103, "grad_norm": 0.007447324987566207, "grad_rms": 6.685480305253421e-06, "grad_to_weight": 0.00021518677588506236, "update_to_weight": 8.994407536066942e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.04810896563666, "grad_norm": 0.0060936000867431546, "grad_rms": 4.901459226325554e-06, "grad_to_weight": 0.00016447803293806896, "update_to_weight": 9.262259724600743e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665758820744294, "grad_norm": 0.00219819581935697, "grad_rms": 3.971461411829867e-06, "grad_to_weight": 8.911932672868966e-05, "update_to_weight": 5.904888267143564e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30663034862235, "grad_norm": 0.00623503661891278, "grad_rms": 5.01522537206434e-06, "grad_to_weight": 0.0001671294501982012, "update_to_weight": 9.368790929154096e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.42807234441328, "grad_norm": 0.0160936747460666, "grad_rms": 7.145033972149472e-06, "grad_to_weight": 0.00025373110282586294, "update_to_weight": 0.00010421990702890436, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76145527658634, "grad_norm": 0.035817898819263734, "grad_rms": 2.6726277652876125e-05, "grad_to_weight": 0.0007056121347211394, "update_to_weight": 9.446514936773014e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.076426925528025, "grad_norm": 0.048924853359872754, "grad_rms": 4.201473315070185e-05, "grad_to_weight": 0.0012207888056185074, "update_to_weight": 0.0001251875264882422, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199254192473227, "grad_norm": 0.1564260921029841, "grad_rms": 0.00024384611064420792, "grad_to_weight": 0.00554717124911539, "update_to_weight": 0.00010266520039529748, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40794701446476, "grad_norm": 0.025660991876927004, "grad_rms": 0.002184408692098131, "grad_to_weight": 0.003051992577115515, "update_to_weight": 8.211929227146533e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3111498109490423, "grad_norm": 0.29887592955540065, "grad_rms": 0.0069188691957045596, "grad_to_weight": 0.22794948911220675, "update_to_weight": 0.000266619577022011, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43850}
{"step": 43900, "lr": 1.3614068110586382e-05, "window_steps": 50, "elapsed_seconds": 1671.406496104, "steps_per_second": 1.735065650851433, "eta_seconds": 3515.717112494621, "images_seen": 44953600, "images_per_second": 1776.7072264718674, "effective_batch_size": 1024, "loss": 0.3642294704914093, "edge_loss": 0.1911281943321228, "diag_loss": 0.17310132086277008, "jvp_rms": 2.096445322036743, "grad_norm": 0.28420281410217285, "grad_norm_postclip": 0.28420281410217285, "clip_fraction": 0.0, "ema_updates": 43900, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3390564481417338, "updates_nfe_2": 15, "loss_nfe_3": 0.3739035129547119, "updates_nfe_3": 8, "loss_nfe_5": 0.35657527670264244, "updates_nfe_5": 8, "loss_nfe_4": 0.3158145010471344, "updates_nfe_4": 5, "loss_nfe_1": 0.40733765065670013, "updates_nfe_1": 14, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.492758346645942, "grad_norm": 0.04830305067161423, "grad_rms": 0.00014065376551795463, "grad_to_weight": 0.002928743006862609, "update_to_weight": 8.13746775015236e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.249539882308014, "grad_norm": 0.03498443368052391, "grad_rms": 0.00010187125374610217, "grad_to_weight": 0.0021529491871098367, "update_to_weight": 0.00013590370361422658, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061645058862217, "grad_norm": 0.0029687930486033873, "grad_rms": 5.3149921501243966e-05, "grad_to_weight": 0.0005865273076399254, "update_to_weight": 3.5939145507920345e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.848605757897431, "grad_norm": 0.014326935973167424, "grad_rms": 4.5846253797378175e-05, "grad_to_weight": 0.0011150576368460317, "update_to_weight": 0.00010426967343793464, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.05317922550093, "grad_norm": 0.022225061388562548, "grad_rms": 0.0005056376545967009, "grad_to_weight": 0.004398233349097041, "update_to_weight": 4.6571035456255755e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.638995482128628, "grad_norm": 0.06491225886367152, "grad_rms": 0.0001564155357021632, "grad_to_weight": 0.003901212602250631, "update_to_weight": 9.77390832392307e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.561971954804317, "grad_norm": 0.030785192131899612, "grad_rms": 4.247279943091612e-05, "grad_to_weight": 0.001253368100433729, "update_to_weight": 0.0001203874866862092, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826709372590962, "grad_norm": 0.022036064236807037, "grad_rms": 2.675126862332519e-05, "grad_to_weight": 0.0007148367336411205, "update_to_weight": 9.468312848053654e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61072055792919, "grad_norm": 0.016277312785421955, "grad_rms": 1.4612180109109381e-05, "grad_to_weight": 0.00047029684800055057, "update_to_weight": 0.00010978117116494936, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.049054593728336, "grad_norm": 0.013240782886886807, "grad_rms": 1.065038015637017e-05, "grad_to_weight": 0.0003573851757374723, "update_to_weight": 0.000113601363134672, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665845427717706, "grad_norm": 0.004246275991019254, "grad_rms": 7.67171017878004e-06, "grad_to_weight": 0.00017215205549969085, "update_to_weight": 7.27138199452254e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30792116220512, "grad_norm": 0.011981622228933843, "grad_rms": 9.637559404024359e-06, "grad_to_weight": 0.0003211549144440633, "update_to_weight": 0.00011402704056724123, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.42941566406217, "grad_norm": 0.023252427025411217, "grad_rms": 1.0323271947079408e-05, "grad_to_weight": 0.00036658743868242153, "update_to_weight": 0.0001281400266603174, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76234263395948, "grad_norm": 0.03051317946205124, "grad_rms": 2.2768049864952465e-05, "grad_to_weight": 0.0006010987255272621, "update_to_weight": 0.00010691291728222036, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.076810711880974, "grad_norm": 0.03458705378200802, "grad_rms": 2.9701996742455224e-05, "grad_to_weight": 0.0008630191167321235, "update_to_weight": 0.00012506179309490988, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199266096805385, "grad_norm": 0.06633762712789919, "grad_rms": 0.00010341096007086995, "grad_to_weight": 0.002352459347706726, "update_to_weight": 0.0001050360076152017, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408576274748402, "grad_norm": 0.009570482843084894, "grad_rms": 0.0008146936022690571, "grad_to_weight": 0.0011381811296432888, "update_to_weight": 6.509328378649619e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3118004524091758, "grad_norm": 0.1436229076606593, "grad_rms": 0.0033248181380450006, "grad_to_weight": 0.10948533170337751, "update_to_weight": 0.0001751812934925257, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43900}
{"step": 43950, "lr": 1.3555863158708318e-05, "window_steps": 50, "elapsed_seconds": 1698.890827269, "steps_per_second": 1.736427057377302, "eta_seconds": 3484.165933890661, "images_seen": 45004800, "images_per_second": 1778.1013067543572, "effective_batch_size": 1024, "loss": 0.3668674826622009, "edge_loss": 0.19129599630832672, "diag_loss": 0.175571471452713, "jvp_rms": 1.9404296875, "grad_norm": 0.32111281156539917, "grad_norm_postclip": 0.3124808371067047, "clip_fraction": 0.03999999910593033, "ema_updates": 43950, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.36455739608832766, "updates_nfe_4": 14, "loss_nfe_1": 0.40461940318346024, "updates_nfe_1": 12, "loss_nfe_5": 0.3482249490916729, "updates_nfe_5": 8, "loss_nfe_2": 0.35916287248784845, "updates_nfe_2": 11, "loss_nfe_3": 0.3295090854167938, "updates_nfe_3": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.490957845856865, "grad_norm": 0.03004791392112499, "grad_rms": 8.749659038511324e-05, "grad_to_weight": 0.0018220842113591438, "update_to_weight": 6.0542416947007814e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.249469454980964, "grad_norm": 0.01476044745723588, "grad_rms": 4.29809812573638e-05, "grad_to_weight": 0.0009083648852739218, "update_to_weight": 6.798860768416473e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061760554057868, "grad_norm": 0.0010774248757372648, "grad_rms": 1.9288999479388578e-05, "grad_to_weight": 0.00021285575724705593, "update_to_weight": 3.686812324548099e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.849007748541142, "grad_norm": 0.01029110551377379, "grad_rms": 3.293157979652524e-05, "grad_to_weight": 0.0008009260882376094, "update_to_weight": 8.01326883233548e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.052830822366608, "grad_norm": 0.032214545408023584, "grad_rms": 0.0007329062853520193, "grad_to_weight": 0.006375544034726888, "update_to_weight": 2.618019571488228e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63912111748995, "grad_norm": 0.0855176107004262, "grad_rms": 0.00020606713005888498, "grad_to_weight": 0.005139550947227358, "update_to_weight": 8.54286928257145e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.562706245158278, "grad_norm": 0.015631732011974944, "grad_rms": 2.1566323694127132e-05, "grad_to_weight": 0.000636401048644883, "update_to_weight": 9.757811967627623e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.827006362750545, "grad_norm": 0.007568834995506879, "grad_rms": 9.188389358215405e-06, "grad_to_weight": 0.0002455261113077977, "update_to_weight": 7.475680107718178e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61086868955058, "grad_norm": 0.00449406775451891, "grad_rms": 4.0343346790253845e-06, "grad_to_weight": 0.00012984556368201533, "update_to_weight": 8.26361846891383e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.05024565539343, "grad_norm": 0.0038367053799321245, "grad_rms": 3.0860993034434837e-06, "grad_to_weight": 0.00010355411447517926, "update_to_weight": 8.835023389585488e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66628711855199, "grad_norm": 0.0012641636554050107, "grad_rms": 2.2839535638583098e-06, "grad_to_weight": 5.125066652022423e-05, "update_to_weight": 6.027975781373541e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30887656453639, "grad_norm": 0.0038570699557648595, "grad_rms": 3.1024797906242422e-06, "grad_to_weight": 0.00010338209860307512, "update_to_weight": 9.204074766412956e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.429979541912985, "grad_norm": 0.01030272227044773, "grad_rms": 4.574051718421957e-06, "grad_to_weight": 0.00016242670019529082, "update_to_weight": 0.0001072278996704354, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.762135825768304, "grad_norm": 0.021523926238475398, "grad_rms": 1.606052972934675e-05, "grad_to_weight": 0.00042401537855602287, "update_to_weight": 9.05400666853831e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.077537154558286, "grad_norm": 0.03330558507283855, "grad_rms": 2.8601521990682367e-05, "grad_to_weight": 0.0008310287367309068, "update_to_weight": 0.00010679232183298427, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19916220428255, "grad_norm": 0.10281618454455886, "grad_rms": 0.00016027586175305033, "grad_to_weight": 0.0036460723123520446, "update_to_weight": 8.619785977556743e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408312236153254, "grad_norm": 0.02012619310602468, "grad_rms": 0.0017132553320815213, "grad_to_weight": 0.0023936067715811036, "update_to_weight": 3.298217095881536e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.311794681861893, "grad_norm": 0.2740027715567174, "grad_rms": 0.00634306462377739, "grad_to_weight": 0.2088762634468163, "update_to_weight": 0.00012055951096721477, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 43950}
{"step": 44000, "lr": 1.3498111433270821e-05, "window_steps": 50, "elapsed_seconds": 1726.370491353, "steps_per_second": 1.7377498138587995, "eta_seconds": 3452.740982706, "images_seen": 45056000, "images_per_second": 1779.4558093914106, "effective_batch_size": 1024, "loss": 0.36284881830215454, "edge_loss": 0.18861538171768188, "diag_loss": 0.17423343658447266, "jvp_rms": 1.8724608421325684, "grad_norm": 0.3744378089904785, "grad_norm_postclip": 0.3374742567539215, "clip_fraction": 0.05999999865889549, "ema_updates": 44000, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.35151649605144153, "updates_nfe_3": 11, "loss_nfe_4": 0.30481454978386563, "updates_nfe_4": 6, "loss_nfe_1": 0.3985885977745056, "updates_nfe_1": 11, "loss_nfe_2": 0.36173719763755796, "updates_nfe_2": 15, "loss_nfe_5": 0.37662006488868166, "updates_nfe_5": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.490690437145133, "grad_norm": 0.018987686199353322, "grad_rms": 5.5290287575600324e-05, "grad_to_weight": 0.0011514185092325623, "update_to_weight": 4.920503595118779e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251920146470866, "grad_norm": 0.027808017995757425, "grad_rms": 8.09742322340076e-05, "grad_to_weight": 0.0017110604620953658, "update_to_weight": 6.087187479706682e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061515806627713, "grad_norm": 0.0018831968656126481, "grad_rms": 3.371463215524141e-05, "grad_to_weight": 0.00037206183632711946, "update_to_weight": 2.9575249129302837e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.849193004397506, "grad_norm": 0.014181159531884042, "grad_rms": 4.53797685881699e-05, "grad_to_weight": 0.001103661492751388, "update_to_weight": 7.836295554778846e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.052402551797527, "grad_norm": 0.03806028357952224, "grad_rms": 0.0008659014337903545, "grad_to_weight": 0.007533105921257459, "update_to_weight": 3.369402753952836e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63910919592706, "grad_norm": 0.07969879229326376, "grad_rms": 0.00019204584017862644, "grad_to_weight": 0.004789847302208494, "update_to_weight": 9.401358831470417e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.562933609801807, "grad_norm": 0.028556214269259098, "grad_rms": 3.93975894634011e-05, "grad_to_weight": 0.0011625734418735626, "update_to_weight": 0.00011104395582214492, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826660863926396, "grad_norm": 0.01590070336653639, "grad_rms": 1.9303083458412546e-05, "grad_to_weight": 0.0005158101111477669, "update_to_weight": 9.433562464500935e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.610274395762886, "grad_norm": 0.011171987507368662, "grad_rms": 1.0029118183475334e-05, "grad_to_weight": 0.0003227939593780388, "update_to_weight": 0.0001117448824732827, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.050560300145094, "grad_norm": 0.00950299279978278, "grad_rms": 7.643844537408007e-06, "grad_to_weight": 0.00025648715492557786, "update_to_weight": 0.00011745601626867195, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66590605241816, "grad_norm": 0.0034416920902334345, "grad_rms": 6.218075390462971e-06, "grad_to_weight": 0.00013953236029195138, "update_to_weight": 7.649298720313564e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30974687550242, "grad_norm": 0.008631964119024083, "grad_rms": 6.943222326740304e-06, "grad_to_weight": 0.0002313594929450414, "update_to_weight": 0.00011483707825657391, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.43144791019968, "grad_norm": 0.019339266855259, "grad_rms": 8.585964415060761e-06, "grad_to_weight": 0.00030488452482809044, "update_to_weight": 0.00012119141447620605, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.762479704040516, "grad_norm": 0.028182070499953624, "grad_rms": 2.10286439418272e-05, "grad_to_weight": 0.0005551752133517313, "update_to_weight": 9.308021365198831e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07778995961323, "grad_norm": 0.03121108933257888, "grad_rms": 2.6802851712306678e-05, "grad_to_weight": 0.0007787627352713459, "update_to_weight": 0.0001127683728156727, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199192506307917, "grad_norm": 0.06979187970105201, "grad_rms": 0.00010879565033463698, "grad_to_weight": 0.0024749602204190836, "update_to_weight": 8.42759156453362e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40807586467283, "grad_norm": 0.009801706697698993, "grad_rms": 0.0008343766838998035, "grad_to_weight": 0.001165749079272894, "update_to_weight": 6.840540528885498e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3116413214018958, "grad_norm": 0.10110778988211862, "grad_rms": 0.0023406086060586834, "grad_to_weight": 0.07708493795701218, "update_to_weight": 0.0002087111238492463, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44000}
{"step": 44050, "lr": 1.3440813565837496e-05, "window_steps": 50, "elapsed_seconds": 1776.794965538, "steps_per_second": 1.7165739768271366, "eta_seconds": 3466.206572115115, "images_seen": 45107200, "images_per_second": 1757.771752270988, "effective_batch_size": 1024, "loss": 0.3880464732646942, "edge_loss": 0.2013336569070816, "diag_loss": 0.1867128312587738, "jvp_rms": 2.509765625, "grad_norm": 0.255149245262146, "grad_norm_postclip": 0.255149245262146, "clip_fraction": 0.0, "ema_updates": 44050, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4084348529577255, "updates_nfe_1": 10, "loss_nfe_3": 0.4066140428185463, "updates_nfe_3": 12, "loss_nfe_2": 0.37275766432285307, "updates_nfe_2": 10, "loss_nfe_5": 0.38649644383362364, "updates_nfe_5": 7, "loss_nfe_4": 0.3641414127566598, "updates_nfe_4": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.494547549931557, "grad_norm": 0.01754811078883792, "grad_rms": 5.10983845918197e-05, "grad_to_weight": 0.0010638734245796719, "update_to_weight": 9.281957500069626e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.24791340269286, "grad_norm": 0.006021834633915056, "grad_rms": 1.7534994266611214e-05, "grad_to_weight": 0.000370622029098027, "update_to_weight": 8.885795627106709e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.060775086216202, "grad_norm": 0.0006711551816665437, "grad_rms": 1.201560520950124e-05, "grad_to_weight": 0.00013261904949985585, "update_to_weight": 4.389267846936912e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.848744704387103, "grad_norm": 0.006881881974021499, "grad_rms": 2.2022050505111483e-05, "grad_to_weight": 0.0005356073400440227, "update_to_weight": 0.00011270650426574219, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.052868192867028, "grad_norm": 0.014095833423966867, "grad_rms": 0.0003206913145242552, "grad_to_weight": 0.002789669725378054, "update_to_weight": 3.218581604963709e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63942007081283, "grad_norm": 0.062260459281177516, "grad_rms": 0.00015002563863908769, "grad_to_weight": 0.0037417445449549323, "update_to_weight": 0.00010429549332086494, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.563320000625424, "grad_norm": 0.01846031687647175, "grad_rms": 2.546878163911462e-05, "grad_to_weight": 0.0007515399740752357, "update_to_weight": 0.00011805273723449206, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.826662843873383, "grad_norm": 0.00847442359900921, "grad_rms": 1.0287752825946112e-05, "grad_to_weight": 0.00027490564392030687, "update_to_weight": 9.81699044674917e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.611082068509866, "grad_norm": 0.005045574018230882, "grad_rms": 4.5294230859938085e-06, "grad_to_weight": 0.00014577914692882389, "update_to_weight": 0.00011243757907378784, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.051571757842645, "grad_norm": 0.004125547034009708, "grad_rms": 3.3184325005966263e-06, "grad_to_weight": 0.00011134607354778303, "update_to_weight": 0.00011365021540279896, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665443321340916, "grad_norm": 0.001452674756909467, "grad_rms": 2.6245349437034844e-06, "grad_to_weight": 5.889514078397248e-05, "update_to_weight": 7.732631238981584e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.30999062437064, "grad_norm": 0.0042314784118215945, "grad_rms": 3.4036396559304706e-06, "grad_to_weight": 0.00011341408402975102, "update_to_weight": 0.0001059900199830866, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.43242552056381, "grad_norm": 0.010824194248548418, "grad_rms": 4.805567208690209e-06, "grad_to_weight": 0.00017064134249508372, "update_to_weight": 0.00010950873652014664, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76279953279025, "grad_norm": 0.019380684243977205, "grad_rms": 1.4461304690734055e-05, "grad_to_weight": 0.00038178911372802136, "update_to_weight": 9.095897796622522e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0783244999908, "grad_norm": 0.023432990723192444, "grad_rms": 2.0123327604396997e-05, "grad_to_weight": 0.0005846798990610953, "update_to_weight": 0.00011214058418899893, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199421934871953, "grad_norm": 0.05219904171283403, "grad_rms": 8.137090896990152e-05, "grad_to_weight": 0.001851067792573566, "update_to_weight": 8.885849415899391e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408802198098032, "grad_norm": 0.00734390825932692, "grad_rms": 0.0006251549867045155, "grad_to_weight": 0.0008733596160685076, "update_to_weight": 5.537421731303477e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3120663698581176, "grad_norm": 0.09528062944219533, "grad_rms": 0.0022057119587234886, "grad_to_weight": 0.07261875742802451, "update_to_weight": 0.00012126575463909234, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44050}
{"step": 44100, "lr": 1.3383970183008648e-05, "window_steps": 50, "elapsed_seconds": 1804.279264129, "steps_per_second": 1.7181375752807855, "eta_seconds": 3433.9508575358386, "images_seen": 45158400, "images_per_second": 1759.3728770875243, "effective_batch_size": 1024, "loss": 0.38807323575019836, "edge_loss": 0.2004777044057846, "diag_loss": 0.1875954270362854, "jvp_rms": 2.5921483039855957, "grad_norm": 0.29970911145210266, "grad_norm_postclip": 0.29970911145210266, "clip_fraction": 0.0, "ema_updates": 44100, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3872564978459302, "updates_nfe_2": 17, "loss_nfe_3": 0.3965134173631668, "updates_nfe_3": 8, "loss_nfe_5": 0.37766096327039933, "updates_nfe_5": 9, "loss_nfe_4": 0.3747054537137349, "updates_nfe_4": 9, "loss_nfe_1": 0.41098469495773315, "updates_nfe_1": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.491431583517183, "grad_norm": 0.030524332252773827, "grad_rms": 8.888387403234609e-05, "grad_to_weight": 0.0018509207098359013, "update_to_weight": 4.517089340735469e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25490930199291, "grad_norm": 0.009120173324244108, "grad_rms": 2.6557053900224524e-05, "grad_to_weight": 0.0005610719293971047, "update_to_weight": 6.343231062272067e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061466818037797, "grad_norm": 0.00100184553631819, "grad_rms": 1.793591225118672e-05, "grad_to_weight": 0.00019793581037573223, "update_to_weight": 2.5544113717923066e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.850728981328745, "grad_norm": 0.010786908973581167, "grad_rms": 3.4518152898723715e-05, "grad_to_weight": 0.0008394005499029533, "update_to_weight": 6.428072808650676e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.05179906091197, "grad_norm": 0.013952197754663892, "grad_rms": 0.0003174234898972258, "grad_to_weight": 0.0027618275363757616, "update_to_weight": 3.1537833445248525e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63862957519714, "grad_norm": 0.07366333045741381, "grad_rms": 0.00017750251642452476, "grad_to_weight": 0.0044272474559576835, "update_to_weight": 7.177894774976201e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56372502113762, "grad_norm": 0.01823230796879523, "grad_rms": 2.5154209082194554e-05, "grad_to_weight": 0.0007422452398040579, "update_to_weight": 8.044740196727035e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8270608106361, "grad_norm": 0.007872848517765983, "grad_rms": 9.557454691828488e-06, "grad_to_weight": 0.00025538758190821924, "update_to_weight": 6.422059209047529e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61202021493942, "grad_norm": 0.004683255997572179, "grad_rms": 4.20416938813635e-06, "grad_to_weight": 0.00013530721317303425, "update_to_weight": 7.373844330830889e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.05318443214429, "grad_norm": 0.0038097990297951347, "grad_rms": 3.0644568628093167e-06, "grad_to_weight": 0.00010281974648554272, "update_to_weight": 7.247167013612086e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665173597082564, "grad_norm": 0.0013128840268832474, "grad_rms": 2.3719762384498913e-06, "grad_to_weight": 5.322825001477133e-05, "update_to_weight": 4.412335561099727e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31094597371025, "grad_norm": 0.004016028873663902, "grad_rms": 3.2303402743533953e-06, "grad_to_weight": 0.00010763674757787288, "update_to_weight": 7.700017948952546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.43397273137143, "grad_norm": 0.010322523222684643, "grad_rms": 4.582842655149954e-06, "grad_to_weight": 0.0001627286259745736, "update_to_weight": 8.269884290070087e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76342235141485, "grad_norm": 0.02146608077702947, "grad_rms": 1.6017367123093277e-05, "grad_to_weight": 0.00042286512182784656, "update_to_weight": 7.212472333565503e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07878441221661, "grad_norm": 0.028096655961083312, "grad_rms": 2.4128299250053428e-05, "grad_to_weight": 0.0007010356320217893, "update_to_weight": 9.599746687668955e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19944574338962, "grad_norm": 0.06942256365053931, "grad_rms": 0.00010821993894719968, "grad_to_weight": 0.002461841423490141, "update_to_weight": 8.606906792302996e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408368492494034, "grad_norm": 0.005802944199838846, "grad_rms": 0.0004939794147741419, "grad_to_weight": 0.0006901391399555106, "update_to_weight": 6.523767937036653e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3118066773017991, "grad_norm": 0.09031809106435552, "grad_rms": 0.002090830998031828, "grad_to_weight": 0.0688501534769796, "update_to_weight": 0.00023070580860391564, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44100}
{"step": 44150, "lr": 1.3327581906414382e-05, "window_steps": 50, "elapsed_seconds": 1831.776131422, "steps_per_second": 1.7196424530079821, "eta_seconds": 3401.869958355143, "images_seen": 45209600, "images_per_second": 1760.9138718801737, "effective_batch_size": 1024, "loss": 0.36789947748184204, "edge_loss": 0.18889985978603363, "diag_loss": 0.17899955809116364, "jvp_rms": 2.0644922256469727, "grad_norm": 0.2145845592021942, "grad_norm_postclip": 0.2145845592021942, "clip_fraction": 0.0, "ema_updates": 44150, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3671387759121982, "updates_nfe_2": 11, "loss_nfe_3": 0.3781572669744492, "updates_nfe_3": 15, "loss_nfe_4": 0.3280802518129349, "updates_nfe_4": 10, "loss_nfe_1": 0.40147948761781055, "updates_nfe_1": 6, "loss_nfe_5": 0.3743008840829134, "updates_nfe_5": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49514699139047, "grad_norm": 0.0023515164386734813, "grad_rms": 6.847386183232603e-06, "grad_to_weight": 0.0001425580772272499, "update_to_weight": 4.94325801204998e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.248717272088935, "grad_norm": 0.0019629307276184138, "grad_rms": 5.715862548049736e-06, "grad_to_weight": 0.0001208052731024016, "update_to_weight": 6.441739120487426e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061445149856401, "grad_norm": 0.0007519443799511239, "grad_rms": 1.3461963873333964e-05, "grad_to_weight": 0.00014856317863534636, "update_to_weight": 2.792895904982171e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.849645444973314, "grad_norm": 0.0035485771512607784, "grad_rms": 1.1355461419034409e-05, "grad_to_weight": 0.00027616148371229616, "update_to_weight": 7.562861469828579e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.052394624014199, "grad_norm": 0.012841420113726111, "grad_rms": 0.00029215242354006956, "grad_to_weight": 0.002541650260787314, "update_to_weight": 3.416866797728694e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639324699784055, "grad_norm": 0.046342982202376606, "grad_rms": 0.00011167016083116712, "grad_to_weight": 0.002785148017634276, "update_to_weight": 7.855084436933613e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56383559306178, "grad_norm": 0.01748647157197358, "grad_rms": 2.4125215676704042e-05, "grad_to_weight": 0.0007118787090772074, "update_to_weight": 8.533812987980647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82721821398259, "grad_norm": 0.009655719565891507, "grad_rms": 1.172181860984119e-05, "grad_to_weight": 0.00031322059288216514, "update_to_weight": 6.64855474639823e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61208722442566, "grad_norm": 0.006504711834298097, "grad_rms": 5.839294368401186e-06, "grad_to_weight": 0.0001879317994353065, "update_to_weight": 7.532181496867586e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.05211536376415, "grad_norm": 0.005353410320073063, "grad_rms": 4.3060788420811486e-06, "grad_to_weight": 0.00014448325736642114, "update_to_weight": 8.049782529381263e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66589120556652, "grad_norm": 0.0018243220101542101, "grad_rms": 3.295986828051906e-06, "grad_to_weight": 7.396132557912616e-05, "update_to_weight": 5.1220972618984817e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.311148818709334, "grad_norm": 0.0055311124042494225, "grad_rms": 4.449015613058028e-06, "grad_to_weight": 0.0001482428866268491, "update_to_weight": 8.830829028061792e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.43435183075098, "grad_norm": 0.0138420478396839, "grad_rms": 6.14538964028908e-06, "grad_to_weight": 0.00021821059788891406, "update_to_weight": 0.00010221929168297383, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76346804051118, "grad_norm": 0.02328827035067639, "grad_rms": 1.737703215333987e-05, "grad_to_weight": 0.0004587604285052287, "update_to_weight": 8.355917690939638e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.079177313182996, "grad_norm": 0.024205093128763296, "grad_rms": 2.0786378677773963e-05, "grad_to_weight": 0.0006039318856178634, "update_to_weight": 0.00010244887608583553, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199706553562418, "grad_norm": 0.04446128359713367, "grad_rms": 6.930884057547417e-05, "grad_to_weight": 0.0015766576688549605, "update_to_weight": 7.871233796642469e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408825334483966, "grad_norm": 0.006403716177024458, "grad_rms": 0.0005451205216817517, "grad_to_weight": 0.0007615470558965347, "update_to_weight": 3.6465690329538483e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.312183205655896, "grad_norm": 0.08405627178408417, "grad_rms": 0.0019458721564423198, "grad_to_weight": 0.06405833531611813, "update_to_weight": 9.602180449642952e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44150}
{"step": 44200, "lr": 1.3271649352707893e-05, "window_steps": 50, "elapsed_seconds": 1859.239670387, "steps_per_second": 1.721133671450718, "eta_seconds": 3369.8719025764376, "images_seen": 45260800, "images_per_second": 1762.4408795655352, "effective_batch_size": 1024, "loss": 0.392577201128006, "edge_loss": 0.2031785398721695, "diag_loss": 0.1893986314535141, "jvp_rms": 2.4482812881469727, "grad_norm": 0.2756043076515198, "grad_norm_postclip": 0.27245306968688965, "clip_fraction": 0.019999999552965164, "ema_updates": 44200, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.36219896127780277, "updates_nfe_2": 12, "loss_nfe_4": 0.39008577380861553, "updates_nfe_4": 7, "loss_nfe_1": 0.39773207902908325, "updates_nfe_1": 12, "loss_nfe_3": 0.33913054317235947, "updates_nfe_3": 10, "loss_nfe_5": 0.487531046072642, "updates_nfe_5": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.493516977737702, "grad_norm": 0.10675864518389049, "grad_rms": 0.0003108707470423551, "grad_to_weight": 0.006472764136841711, "update_to_weight": 4.988549859373828e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25064414288227, "grad_norm": 0.04260579755284949, "grad_rms": 0.00012406392091971093, "grad_to_weight": 0.002621791307362465, "update_to_weight": 4.750813714261373e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061782597702735, "grad_norm": 0.003487870431990339, "grad_rms": 6.24428973767656e-05, "grad_to_weight": 0.0006890597066680208, "update_to_weight": 3.098726927452564e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.850068777174892, "grad_norm": 0.02199043839195941, "grad_rms": 7.03694929272787e-05, "grad_to_weight": 0.0017113090033432524, "update_to_weight": 6.390979920171474e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.052324028442644, "grad_norm": 0.022899457062519786, "grad_rms": 0.000520980687440938, "grad_to_weight": 0.004532460098284401, "update_to_weight": 2.6858595831550923e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639352210713838, "grad_norm": 0.0777056226319348, "grad_rms": 0.00018724300777409045, "grad_to_weight": 0.004669990853484144, "update_to_weight": 7.805606517208382e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56424681686144, "grad_norm": 0.029108762810413056, "grad_rms": 4.0159913221645887e-05, "grad_to_weight": 0.0011850053057776704, "update_to_weight": 9.294225441270121e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.827497380205266, "grad_norm": 0.02429814117090091, "grad_rms": 2.9497377323149046e-05, "grad_to_weight": 0.0007881970070818354, "update_to_weight": 8.697132905086569e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61193557119242, "grad_norm": 0.018187977316744344, "grad_rms": 1.6327387934124218e-05, "grad_to_weight": 0.0005254828144278136, "update_to_weight": 0.00010493157532969067, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.052388811000604, "grad_norm": 0.014979408359050551, "grad_rms": 1.2048864096955817e-05, "grad_to_weight": 0.00040427645395438755, "update_to_weight": 0.00011257744048999474, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665974100373827, "grad_norm": 0.005405896784837961, "grad_rms": 9.766787057032728e-06, "grad_to_weight": 0.0002191641312376158, "update_to_weight": 7.5228526637215e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.312475463245704, "grad_norm": 0.01372920728141354, "grad_rms": 1.1043250088895567e-05, "grad_to_weight": 0.0003679521959066309, "update_to_weight": 0.00011183859948398282, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.43460969400892, "grad_norm": 0.02806498263099343, "grad_rms": 1.245987989009425e-05, "grad_to_weight": 0.0004424238245710216, "update_to_weight": 0.00011627204487728552, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76345842175748, "grad_norm": 0.03439779176264053, "grad_rms": 2.5666634939504346e-05, "grad_to_weight": 0.000677609304646932, "update_to_weight": 8.538404546242956e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07955650464929, "grad_norm": 0.036674736788771534, "grad_rms": 3.149481651418222e-05, "grad_to_weight": 0.0009150484682762697, "update_to_weight": 9.904438708141485e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199548552666915, "grad_norm": 0.08189924722817243, "grad_rms": 0.0001276693205896262, "grad_to_weight": 0.00290427511898686, "update_to_weight": 8.294876641409647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40895643946851, "grad_norm": 0.003880352548058347, "grad_rms": 0.0003303175448181019, "grad_to_weight": 0.00046145470915337576, "update_to_weight": 3.409751725016755e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3124594455084817, "grad_norm": 0.06416745891845722, "grad_rms": 0.0014854533636682722, "grad_to_weight": 0.04889100317579492, "update_to_weight": 0.00011329719674465094, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44200}
{"step": 44250, "lr": 1.3216173133558612e-05, "window_steps": 50, "elapsed_seconds": 1886.71662513, "steps_per_second": 1.7225692277853681, "eta_seconds": 3338.037105999231, "images_seen": 45312000, "images_per_second": 1763.910889252217, "effective_batch_size": 1024, "loss": 0.3785541355609894, "edge_loss": 0.19805966317653656, "diag_loss": 0.1804945319890976, "jvp_rms": 2.489023447036743, "grad_norm": 0.20782841742038727, "grad_norm_postclip": 0.20782841742038727, "clip_fraction": 0.0, "ema_updates": 44250, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.401657106479009, "updates_nfe_1": 15, "loss_nfe_5": 0.3429406713694334, "updates_nfe_5": 8, "loss_nfe_2": 0.37878626585006714, "updates_nfe_2": 10, "loss_nfe_3": 0.39090717732906344, "updates_nfe_3": 10, "loss_nfe_4": 0.351770213672093, "updates_nfe_4": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.494576227411088, "grad_norm": 0.021965972228142446, "grad_rms": 6.396276558504582e-05, "grad_to_weight": 0.0013317087947757554, "update_to_weight": 5.753701354206924e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25109202310292, "grad_norm": 0.011955815651118171, "grad_rms": 3.481416738252713e-05, "grad_to_weight": 0.0007356930619875584, "update_to_weight": 4.9037126282284046e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0620874307437775, "grad_norm": 0.0009361626568901286, "grad_rms": 1.6760000077982382e-05, "grad_to_weight": 0.00018493609004153002, "update_to_weight": 2.7215262363395948e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.850442224158876, "grad_norm": 0.008631364368651936, "grad_rms": 2.762040133382253e-05, "grad_to_weight": 0.0006716783919252943, "update_to_weight": 6.51571782961132e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.052058814339932, "grad_norm": 0.017536549608524885, "grad_rms": 0.0003989703181803773, "grad_to_weight": 0.003471168933890587, "update_to_weight": 2.7207743731090994e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639464088026354, "grad_norm": 0.06947259799996659, "grad_rms": 0.00016740433660778598, "grad_to_weight": 0.004175170404073206, "update_to_weight": 8.180183303904561e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.5649127114964, "grad_norm": 0.021413433333048595, "grad_rms": 2.9543049631951294e-05, "grad_to_weight": 0.0008717080978279678, "update_to_weight": 8.410998296310236e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82760429425935, "grad_norm": 0.010502990817136406, "grad_rms": 1.27503861705136e-05, "grad_to_weight": 0.00034070084450552816, "update_to_weight": 6.849589581909306e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61239758245952, "grad_norm": 0.006605247179569844, "grad_rms": 5.929545172806564e-06, "grad_to_weight": 0.00019083471937572958, "update_to_weight": 7.848663313268904e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.052584835187275, "grad_norm": 0.005302653205036948, "grad_rms": 4.2652517569009665e-06, "grad_to_weight": 0.00014311155965565032, "update_to_weight": 8.78468621333466e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665787277354816, "grad_norm": 0.0019285297902891065, "grad_rms": 3.484258124891719e-06, "grad_to_weight": 7.818642756477724e-05, "update_to_weight": 5.955032591073485e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31290076470285, "grad_norm": 0.005579178635916871, "grad_rms": 4.487678254407542e-06, "grad_to_weight": 0.00014952411958264703, "update_to_weight": 9.179275504781413e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.43500610861482, "grad_norm": 0.013423466004025116, "grad_rms": 5.9595538083180336e-06, "grad_to_weight": 0.0002116097534702079, "update_to_weight": 9.898170529396457e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.763780649014414, "grad_norm": 0.023079175711214733, "grad_rms": 1.7221011795524386e-05, "grad_to_weight": 0.0004546386304595069, "update_to_weight": 8.017619320847351e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.07980472844538, "grad_norm": 0.03097935542519653, "grad_rms": 2.660384777847733e-05, "grad_to_weight": 0.0007729417754176309, "update_to_weight": 9.736208895009361e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199696813806792, "grad_norm": 0.05463999575322693, "grad_rms": 8.517601041435369e-05, "grad_to_weight": 0.0019376093336746357, "update_to_weight": 7.744438390357234e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.408986380112527, "grad_norm": 0.006290930826245677, "grad_rms": 0.0005355195950393164, "grad_to_weight": 0.0007481199923362813, "update_to_weight": 3.937487198799498e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3123955003826948, "grad_norm": 0.08446330641671634, "grad_rms": 0.001955294860323134, "grad_to_weight": 0.0643581194785313, "update_to_weight": 7.324787734812901e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44250}
{"step": 44300, "lr": 1.3161153855645575e-05, "window_steps": 50, "elapsed_seconds": 1914.166094963, "steps_per_second": 1.723988325090351, "eta_seconds": 3306.2868912997274, "images_seen": 45363200, "images_per_second": 1765.3640448925194, "effective_batch_size": 1024, "loss": 0.37866678833961487, "edge_loss": 0.19437524676322937, "diag_loss": 0.18429160118103027, "jvp_rms": 2.2744140625, "grad_norm": 0.4024190902709961, "grad_norm_postclip": 0.37832579016685486, "clip_fraction": 0.07999999821186066, "ema_updates": 44300, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.36536728259589935, "updates_nfe_5": 18, "loss_nfe_2": 0.37941851913928987, "updates_nfe_2": 10, "loss_nfe_1": 0.39964216521808077, "updates_nfe_1": 7, "loss_nfe_3": 0.39644045721400867, "updates_nfe_3": 11, "loss_nfe_4": 0.3510511741042137, "updates_nfe_4": 4, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.493358778279894, "grad_norm": 0.06760490829377408, "grad_rms": 0.00019685888958983117, "grad_to_weight": 0.004098916976377364, "update_to_weight": 5.412766483904081e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251970846530718, "grad_norm": 0.0553574440966872, "grad_rms": 0.00016119547012844016, "grad_to_weight": 0.003406198830863905, "update_to_weight": 4.460647451572692e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06238941966662, "grad_norm": 0.0035084124035854067, "grad_rms": 6.28106576617987e-05, "grad_to_weight": 0.0006930348720222418, "update_to_weight": 3.0268467002944178e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.84989481485859, "grad_norm": 0.02974384014098846, "grad_rms": 9.51804102821662e-05, "grad_to_weight": 0.0023147146781773704, "update_to_weight": 6.328044311346239e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.051944796311178, "grad_norm": 0.07364991189345395, "grad_rms": 0.0016755935140059612, "grad_to_weight": 0.014578526659125717, "update_to_weight": 2.9214262163910466e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639394394051163, "grad_norm": 0.1511411086029074, "grad_rms": 0.00036419649974580215, "grad_to_weight": 0.009083329899130382, "update_to_weight": 7.634973206211126e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.5644145345616, "grad_norm": 0.04217675588806031, "grad_rms": 5.818917373668375e-05, "grad_to_weight": 0.0017169860013849924, "update_to_weight": 7.716628449108431e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8278567287491, "grad_norm": 0.027618120518656097, "grad_rms": 3.352775490787875e-05, "grad_to_weight": 0.0008958819538336668, "update_to_weight": 6.748953555129072e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61298831307314, "grad_norm": 0.018708164480218142, "grad_rms": 1.6794361114730138e-05, "grad_to_weight": 0.000540495501601986, "update_to_weight": 7.714958435682468e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.053293149234726, "grad_norm": 0.015168226428499347, "grad_rms": 1.2200742142023166e-05, "grad_to_weight": 0.0004093624382428913, "update_to_weight": 8.202494985058713e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666184429231752, "grad_norm": 0.005678249918628427, "grad_rms": 1.0258845112877898e-05, "grad_to_weight": 0.0002302038215484664, "update_to_weight": 5.4401994020982384e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31346182805544, "grad_norm": 0.013124701968070545, "grad_rms": 1.0557009097811516e-05, "grad_to_weight": 0.0003517417394438132, "update_to_weight": 8.16031296797106e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.436330038472434, "grad_norm": 0.025172503676413393, "grad_rms": 1.1175719453134204e-05, "grad_to_weight": 0.00039681525808865273, "update_to_weight": 8.913908839261803e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.764187037190645, "grad_norm": 0.03501849704062064, "grad_rms": 2.6129787222211114e-05, "grad_to_weight": 0.0006898268067400653, "update_to_weight": 7.000483330896136e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.080037722388994, "grad_norm": 0.04466483445840777, "grad_rms": 3.835639704807806e-05, "grad_to_weight": 0.0011143910284659656, "update_to_weight": 8.605697906243811e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19953556625257, "grad_norm": 0.12517406592798633, "grad_rms": 0.00019512863052262395, "grad_to_weight": 0.004438869769110196, "update_to_weight": 7.138179163251433e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.409168743551263, "grad_norm": 0.011176138569487331, "grad_rms": 0.0009513760945940829, "grad_to_weight": 0.0013290420147720277, "update_to_weight": 4.8918197657087115e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3126836375787363, "grad_norm": 0.10504356085244519, "grad_rms": 0.0024317202742631084, "grad_to_weight": 0.08002199299611865, "update_to_weight": 7.320146758752513e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44300}
{"step": 44350, "lr": 1.3106592120650801e-05, "window_steps": 50, "elapsed_seconds": 1941.629490618, "steps_per_second": 1.7253549228559206, "eta_seconds": 3274.6885438781196, "images_seen": 45414400, "images_per_second": 1766.7634410044627, "effective_batch_size": 1024, "loss": 0.35423240065574646, "edge_loss": 0.18232682347297668, "diag_loss": 0.17190547287464142, "jvp_rms": 2.002070188522339, "grad_norm": 0.13415543735027313, "grad_norm_postclip": 0.13415543735027313, "clip_fraction": 0.0, "ema_updates": 44350, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.35510447892275726, "updates_nfe_3": 11, "loss_nfe_2": 0.374762412160635, "updates_nfe_2": 8, "loss_nfe_5": 0.37267240285873415, "updates_nfe_5": 10, "loss_nfe_4": 0.2968711828192075, "updates_nfe_4": 12, "loss_nfe_1": 0.39091025127304924, "updates_nfe_1": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.494777893439675, "grad_norm": 0.027892665734950497, "grad_rms": 8.122071818250451e-05, "grad_to_weight": 0.0016909997767259421, "update_to_weight": 2.548102875121877e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25252102589828, "grad_norm": 0.0101677129405687, "grad_rms": 2.960738694372054e-05, "grad_to_weight": 0.0006256083548125562, "update_to_weight": 4.739345984265032e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062613780308964, "grad_norm": 0.0024636554000176756, "grad_rms": 4.410650691150566e-05, "grad_to_weight": 0.0004866370430231244, "update_to_weight": 2.1743003663930814e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.851048384473467, "grad_norm": 0.009646551052997662, "grad_rms": 3.086900288194149e-05, "grad_to_weight": 0.0007506431198759275, "update_to_weight": 4.6186208266891805e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.051609334306749, "grad_norm": 0.01874692442748246, "grad_rms": 0.00042650729879612344, "grad_to_weight": 0.0037110796157904338, "update_to_weight": 2.089915705838567e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639610811230494, "grad_norm": 0.07904591486770571, "grad_rms": 0.00019047263699552684, "grad_to_weight": 0.004750466568265865, "update_to_weight": 5.8522995544345684e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.564674184149837, "grad_norm": 0.014225001500099033, "grad_rms": 1.9625527527312085e-05, "grad_to_weight": 0.0005790836627207375, "update_to_weight": 7.41607233421451e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828106191390614, "grad_norm": 0.006767673168798239, "grad_rms": 8.215797564748108e-06, "grad_to_weight": 0.0002195293193419793, "update_to_weight": 6.401232319929844e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61254218013862, "grad_norm": 0.003392011214050637, "grad_rms": 3.0450160567177315e-06, "grad_to_weight": 9.799948228006905e-05, "update_to_weight": 8.220006413001203e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.053634123494724, "grad_norm": 0.002866563677050506, "grad_rms": 2.3057543623999773e-06, "grad_to_weight": 7.736255147056936e-05, "update_to_weight": 8.748242605326354e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666242578658363, "grad_norm": 0.0010136204312470197, "grad_rms": 1.831299283481245e-06, "grad_to_weight": 4.109342669499329e-05, "update_to_weight": 5.681571005566694e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31368919535011, "grad_norm": 0.002920767372056275, "grad_rms": 2.349353744900472e-06, "grad_to_weight": 7.827602778071727e-05, "update_to_weight": 8.950829470281509e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.437271012517755, "grad_norm": 0.007960305294464261, "grad_rms": 3.5340997413613993e-06, "grad_to_weight": 0.00012548309798656842, "update_to_weight": 9.745966986293588e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7644395252474, "grad_norm": 0.018460819004576925, "grad_rms": 1.377492791817421e-05, "grad_to_weight": 0.00036365651186586124, "update_to_weight": 7.318962551759599e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.080544822172605, "grad_norm": 0.025884773830785843, "grad_rms": 2.2228822172796107e-05, "grad_to_weight": 0.0006458189115350137, "update_to_weight": 7.430553312757644e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199786635870026, "grad_norm": 0.06851252022497417, "grad_rms": 0.00010680131020785041, "grad_to_weight": 0.0024295403759483233, "update_to_weight": 5.445222572286228e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.409184620777934, "grad_norm": 0.0040269141651001475, "grad_rms": 0.00034279369818464143, "grad_to_weight": 0.00047887094251090634, "update_to_weight": 2.285154227586516e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.312680322884748, "grad_norm": 0.0723042578146031, "grad_rms": 0.0016738173022361358, "grad_to_weight": 0.055081390765199544, "update_to_weight": 5.8258654318815394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44350}
{"step": 44400, "lr": 1.3052488525252667e-05, "window_steps": 50, "elapsed_seconds": 1969.0916436470002, "steps_per_second": 1.7266844897593396, "eta_seconds": 3243.209766006824, "images_seen": 45465600, "images_per_second": 1768.1249175135638, "effective_batch_size": 1024, "loss": 0.38162761926651, "edge_loss": 0.199099600315094, "diag_loss": 0.18252798914909363, "jvp_rms": 2.382890462875366, "grad_norm": 0.19622741639614105, "grad_norm_postclip": 0.19622741639614105, "clip_fraction": 0.0, "ema_updates": 44400, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.4392268180847168, "updates_nfe_5": 10, "loss_nfe_1": 0.4003605564435323, "updates_nfe_1": 15, "loss_nfe_4": 0.3395274102687836, "updates_nfe_4": 10, "loss_nfe_3": 0.3676034552710397, "updates_nfe_3": 7, "loss_nfe_2": 0.33940060809254646, "updates_nfe_2": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.495381951698537, "grad_norm": 0.005685150085585514, "grad_rms": 1.6554601747798916e-05, "grad_to_weight": 0.00034465101215798834, "update_to_weight": 5.6220105682659354e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251664765986764, "grad_norm": 0.007661944809245057, "grad_rms": 2.2310834898143897e-05, "grad_to_weight": 0.00047145599663616007, "update_to_weight": 5.114090694962354e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0627980085657285, "grad_norm": 0.0023290865836609015, "grad_rms": 4.169733863713149e-05, "grad_to_weight": 0.0004600394050326181, "update_to_weight": 2.7331371470831434e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.851255576692145, "grad_norm": 0.007737213742540058, "grad_rms": 2.475911566781652e-05, "grad_to_weight": 0.0006020589736440042, "update_to_weight": 5.924103118499753e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.051618584823887, "grad_norm": 0.053047017229308074, "grad_rms": 0.0012068614302672542, "grad_to_weight": 0.010500994154363189, "update_to_weight": 3.74962090403439e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63967408521301, "grad_norm": 0.08547596104274911, "grad_rms": 0.00020596676914661023, "grad_to_weight": 0.005136877117004838, "update_to_weight": 7.07339729350178e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.564383475814633, "grad_norm": 0.029802699277402023, "grad_rms": 4.111730287359094e-05, "grad_to_weight": 0.001213248413368277, "update_to_weight": 8.254791193284344e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82791315499944, "grad_norm": 0.016775023550584748, "grad_rms": 2.036448779336659e-05, "grad_to_weight": 0.0005441504738332994, "update_to_weight": 7.384099754255546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61298831307314, "grad_norm": 0.011702952813965523, "grad_rms": 1.050576692727968e-05, "grad_to_weight": 0.00033810870960093844, "update_to_weight": 8.940451294664261e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.05443795377704, "grad_norm": 0.009782568567770832, "grad_rms": 7.868724609607573e-06, "grad_to_weight": 0.0002640053150981251, "update_to_weight": 9.466507508075302e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666572914289034, "grad_norm": 0.0038419860858915087, "grad_rms": 6.941283097048582e-06, "grad_to_weight": 0.00015575678466731366, "update_to_weight": 6.204711999534062e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31429441089316, "grad_norm": 0.009033767740543524, "grad_rms": 7.2664178170633455e-06, "grad_to_weight": 0.00024209938532044963, "update_to_weight": 9.719252394462457e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.438746908927826, "grad_norm": 0.020680319341825466, "grad_rms": 9.181345254187933e-06, "grad_to_weight": 0.00032598877420315334, "update_to_weight": 0.00010555137898372632, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76506713297664, "grad_norm": 0.03223598588093281, "grad_rms": 2.405355806646713e-05, "grad_to_weight": 0.0006350033143163625, "update_to_weight": 8.264323669362832e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.080876793931445, "grad_norm": 0.033284343930501156, "grad_rms": 2.8583280935966118e-05, "grad_to_weight": 0.0008304295363004799, "update_to_weight": 9.500402058499259e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199770402988214, "grad_norm": 0.0729598356791366, "grad_rms": 0.00011373404477742224, "grad_to_weight": 0.0025872492802780176, "update_to_weight": 7.848166258585567e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.40956566522519, "grad_norm": 0.008721137578641248, "grad_rms": 0.0007423925319463902, "grad_to_weight": 0.0010370497033758183, "update_to_weight": 4.5980818908643e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.313096274307284, "grad_norm": 0.08579387909697556, "grad_rms": 0.00198609713451024, "grad_to_weight": 0.06533708211321793, "update_to_weight": 0.00012783530416534485, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44400}
{"step": 44450, "lr": 1.2998843661119447e-05, "window_steps": 50, "elapsed_seconds": 1996.580899254, "steps_per_second": 1.7279540244470202, "eta_seconds": 3211.8910118433914, "images_seen": 45516800, "images_per_second": 1769.4249210337487, "effective_batch_size": 1024, "loss": 0.37519633769989014, "edge_loss": 0.19474951922893524, "diag_loss": 0.1804467737674713, "jvp_rms": 2.2581639289855957, "grad_norm": 0.2285265177488327, "grad_norm_postclip": 0.2285265177488327, "clip_fraction": 0.0, "ema_updates": 44450, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.40016378462314606, "updates_nfe_3": 6, "loss_nfe_4": 0.37732865413029987, "updates_nfe_4": 9, "loss_nfe_1": 0.4001461457122456, "updates_nfe_1": 11, "loss_nfe_5": 0.35788105002471376, "updates_nfe_5": 14, "loss_nfe_2": 0.3550930440425873, "updates_nfe_2": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.494958280688092, "grad_norm": 0.005937960753476252, "grad_rms": 1.7290761719219474e-05, "grad_to_weight": 0.0003599864062965395, "update_to_weight": 4.264993344098438e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251102351432003, "grad_norm": 0.005961309144252044, "grad_rms": 1.7358749952586477e-05, "grad_to_weight": 0.00036682490918696044, "update_to_weight": 5.683226877900942e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063057197707316, "grad_norm": 0.0007346465691468562, "grad_rms": 1.3152283383202562e-05, "grad_to_weight": 0.00014509940149985334, "update_to_weight": 2.5708032616706033e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.852717107138584, "grad_norm": 0.005172102814244136, "grad_rms": 1.6550750190555038e-05, "grad_to_weight": 0.00040241318400849855, "update_to_weight": 6.817391932757384e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0513182175411675, "grad_norm": 0.01653592812900274, "grad_rms": 0.00037620539127200753, "grad_to_weight": 0.0032735866989294414, "update_to_weight": 2.8195784868665557e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639882245747526, "grad_norm": 0.059523626700472496, "grad_rms": 0.0001434308421902818, "grad_to_weight": 0.003577166341768092, "update_to_weight": 8.11131235579782e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.565290375058414, "grad_norm": 0.019282720766667667, "grad_rms": 2.6603411409487308e-05, "grad_to_weight": 0.000784957982269396, "update_to_weight": 0.00010014382706636676, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82808441298646, "grad_norm": 0.010578038242156345, "grad_rms": 1.2841491996155726e-05, "grad_to_weight": 0.0003431299233662505, "update_to_weight": 9.07906356892388e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61311351140092, "grad_norm": 0.006903977128722417, "grad_rms": 6.197715716590172e-06, "grad_to_weight": 0.00019946131475425853, "update_to_weight": 0.00011038817710740963, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.05533400632509, "grad_norm": 0.005691210952685183, "grad_rms": 4.57779277207425e-06, "grad_to_weight": 0.0001535868210421132, "update_to_weight": 0.00011336042866040067, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666243815878715, "grad_norm": 0.0019611108303791078, "grad_rms": 3.543122008780177e-06, "grad_to_weight": 7.950585606052662e-05, "update_to_weight": 7.327692830825167e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3145806580685, "grad_norm": 0.0058359329733765584, "grad_rms": 4.6942016393239134e-06, "grad_to_weight": 0.0001563981926221824, "update_to_weight": 0.00011207376426104596, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.43978405678786, "grad_norm": 0.013945834624409443, "grad_rms": 6.191467376693279e-06, "grad_to_weight": 0.00021982790187188985, "update_to_weight": 0.00011859688830466468, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.765108011371844, "grad_norm": 0.023071055688636426, "grad_rms": 1.7214952870095217e-05, "grad_to_weight": 0.0004544667901320785, "update_to_weight": 9.003509531204615e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08143261205087, "grad_norm": 0.02746472886354054, "grad_rms": 2.358562519891906e-05, "grad_to_weight": 0.0006852232336446727, "update_to_weight": 9.85060071442786e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199862389194937, "grad_norm": 0.06905697730366114, "grad_rms": 0.00010765004163919704, "grad_to_weight": 0.0024488409322919614, "update_to_weight": 7.779093586821864e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.409533004949093, "grad_norm": 0.010484366318948538, "grad_rms": 0.0008924885299872004, "grad_to_weight": 0.001246723963480301, "update_to_weight": 6.722042427404403e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3131241449700994, "grad_norm": 0.1369427956919959, "grad_rms": 0.003170175972673583, "grad_to_weight": 0.1042877752393428, "update_to_weight": 0.00017334911865361224, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44450}
{"step": 44500, "lr": 1.2945658114902764e-05, "window_steps": 50, "elapsed_seconds": 2024.0587943230003, "steps_per_second": 1.7291987810910736, "eta_seconds": 3180.6638196504286, "images_seen": 45568000, "images_per_second": 1770.6995518372594, "effective_batch_size": 1024, "loss": 0.35033729672431946, "edge_loss": 0.18125231564044952, "diag_loss": 0.16908495128154755, "jvp_rms": 1.7983983755111694, "grad_norm": 0.2618657052516937, "grad_norm_postclip": 0.2618657052516937, "clip_fraction": 0.0, "ema_updates": 44500, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.346620563003752, "updates_nfe_2": 9, "loss_nfe_5": 0.3661145567893982, "updates_nfe_5": 7, "loss_nfe_3": 0.32982302606105807, "updates_nfe_3": 10, "loss_nfe_1": 0.40465455916192794, "updates_nfe_1": 9, "loss_nfe_4": 0.3262904683748881, "updates_nfe_4": 15, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.494487419570667, "grad_norm": 0.017464590653219692, "grad_rms": 5.0855182114793176e-05, "grad_to_weight": 0.0010588137848101906, "update_to_weight": 3.5930193882218457e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25184691277185, "grad_norm": 0.008481239688051971, "grad_rms": 2.4696541559969482e-05, "grad_to_weight": 0.0005218631293768103, "update_to_weight": 5.87163199333461e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062611896549832, "grad_norm": 0.0016839905089590258, "grad_rms": 3.014826628016983e-05, "grad_to_weight": 0.0003326327483460987, "update_to_weight": 3.0393862809256674e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.852112212787365, "grad_norm": 0.008779316704295916, "grad_rms": 2.809384941389719e-05, "grad_to_weight": 0.0006831030229848778, "update_to_weight": 5.283623132949302e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.051482279571531, "grad_norm": 0.018877769352415197, "grad_rms": 0.0004294841239126974, "grad_to_weight": 0.0037370752400256698, "update_to_weight": 2.9396491214098037e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639856569699038, "grad_norm": 0.05349357796484905, "grad_rms": 0.00012890056208905186, "grad_to_weight": 0.0032147860013565354, "update_to_weight": 6.348107179538203e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.565357459425968, "grad_norm": 0.02583888395953289, "grad_rms": 3.564862389781154e-05, "grad_to_weight": 0.0010518423760863393, "update_to_weight": 7.038232696113023e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.827773574086667, "grad_norm": 0.014093259009717292, "grad_rms": 1.710888811611457e-05, "grad_to_weight": 0.0004571611042830501, "update_to_weight": 6.531088255554572e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.613238709275855, "grad_norm": 0.010241054814722731, "grad_rms": 9.193417822259484e-06, "grad_to_weight": 0.0002958710365343037, "update_to_weight": 7.816477841107555e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.05521211809734, "grad_norm": 0.008522962982272616, "grad_rms": 6.8555459745338696e-06, "grad_to_weight": 0.00023000712976920456, "update_to_weight": 8.349917356591845e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666000082271015, "grad_norm": 0.0031498431982568834, "grad_rms": 5.690794516591946e-06, "grad_to_weight": 0.00012769979679521978, "update_to_weight": 5.506967249928854e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.314960136652964, "grad_norm": 0.008947461277624958, "grad_rms": 7.1969961938944676e-06, "grad_to_weight": 0.00023978214755845958, "update_to_weight": 8.825207753732619e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.44079232528055, "grad_norm": 0.02020734619729029, "grad_rms": 8.971361565630777e-06, "grad_to_weight": 0.00031852291651215486, "update_to_weight": 0.00010021539041795937, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76529797524708, "grad_norm": 0.027039176985865596, "grad_rms": 2.0175849936815385e-05, "grad_to_weight": 0.0005326311095238675, "update_to_weight": 7.830992782124508e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08202496863073, "grad_norm": 0.027744655903103387, "grad_rms": 2.382601549262907e-05, "grad_to_weight": 0.0006921969617257885, "update_to_weight": 8.856724086733814e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.19982667693855, "grad_norm": 0.04513660194847349, "grad_rms": 7.036156618670059e-05, "grad_to_weight": 0.001600598559188508, "update_to_weight": 6.743094120806061e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.409698573095948, "grad_norm": 0.005263035045996467, "grad_rms": 0.00044801929545166635, "grad_to_weight": 0.0006258292137644278, "update_to_weight": 3.7260136028406416e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3133930619369218, "grad_norm": 0.06031537020109522, "grad_rms": 0.0013962789092204885, "grad_to_weight": 0.04592332025276982, "update_to_weight": 0.00010465796649793789, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44500}
{"step": 44550, "lr": 1.2892932468231237e-05, "window_steps": 50, "elapsed_seconds": 2053.74081268, "steps_per_second": 1.7285530764553867, "eta_seconds": 3152.9260363678873, "images_seen": 45619200, "images_per_second": 1770.038350290316, "effective_batch_size": 1024, "loss": 0.3563636839389801, "edge_loss": 0.18510814011096954, "diag_loss": 0.17125548422336578, "jvp_rms": 1.8358592987060547, "grad_norm": 0.23035244643688202, "grad_norm_postclip": 0.23035244643688202, "clip_fraction": 0.0, "ema_updates": 44550, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3398265739281972, "updates_nfe_4": 15, "loss_nfe_3": 0.32820237303773564, "updates_nfe_3": 12, "loss_nfe_5": 0.35817989706993103, "updates_nfe_5": 6, "loss_nfe_1": 0.3974505538741748, "updates_nfe_1": 12, "loss_nfe_2": 0.37277408242225646, "updates_nfe_2": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.496606651243745, "grad_norm": 0.0761149695162445, "grad_rms": 0.00022163935664286152, "grad_to_weight": 0.004613977354579518, "update_to_weight": 4.10607717708287e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251813112491767, "grad_norm": 0.03560019701519106, "grad_rms": 0.00010366429643149354, "grad_to_weight": 0.0021905369430951296, "update_to_weight": 6.74087531395664e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062480031668759, "grad_norm": 0.004080986018249516, "grad_rms": 7.306136971038597e-05, "grad_to_weight": 0.0008061238746070252, "update_to_weight": 2.91567087232697e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.852030291638263, "grad_norm": 0.015225156912766786, "grad_rms": 4.872056448321616e-05, "grad_to_weight": 0.0011846499399143587, "update_to_weight": 6.595008325348465e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.051530798621724, "grad_norm": 0.02394316246296526, "grad_rms": 0.0005447258074900886, "grad_to_weight": 0.004739783526510021, "update_to_weight": 2.1046573795607404e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640058309012733, "grad_norm": 0.043296930254137264, "grad_rms": 0.00010433025530945183, "grad_to_weight": 0.0026019698639328927, "update_to_weight": 6.282421274018801e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56488538035543, "grad_norm": 0.021018536073140925, "grad_rms": 2.8998229510510595e-05, "grad_to_weight": 0.0008556333867508893, "update_to_weight": 8.467679808992237e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82776565458331, "grad_norm": 0.01223160118776543, "grad_rms": 1.4848878890122005e-05, "grad_to_weight": 0.0003967722255585169, "update_to_weight": 7.262130395319254e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61400046686763, "grad_norm": 0.00846084417437957, "grad_rms": 7.595318747076485e-06, "grad_to_weight": 0.00024443416132955375, "update_to_weight": 8.866191761793766e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.056567690280964, "grad_norm": 0.006901384362607204, "grad_rms": 5.551210052676662e-06, "grad_to_weight": 0.00018623916872952237, "update_to_weight": 9.333094098068487e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666315574553032, "grad_norm": 0.0030563866260584246, "grad_rms": 5.5219473343256075e-06, "grad_to_weight": 0.00012390932958027755, "update_to_weight": 6.1576319765455e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31592517256372, "grad_norm": 0.0071225569669561214, "grad_rms": 5.729112850163305e-06, "grad_to_weight": 0.00019087177750567826, "update_to_weight": 9.3782347264992e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.44249518510838, "grad_norm": 0.01797168781531358, "grad_rms": 7.978806705327763e-06, "grad_to_weight": 0.00028327523630457724, "update_to_weight": 0.00010539141955235471, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76557931284494, "grad_norm": 0.028144477599977694, "grad_rms": 2.1000593209772725e-05, "grad_to_weight": 0.0005544007963848144, "update_to_weight": 8.013956147985901e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08218942597432, "grad_norm": 0.03688621082105669, "grad_rms": 3.167642206687359e-05, "grad_to_weight": 0.000920264370517481, "update_to_weight": 9.678362398170486e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.199710882341616, "grad_norm": 0.06823420645516319, "grad_rms": 0.00010636745848020861, "grad_to_weight": 0.0024196775186759376, "update_to_weight": 8.253277888805944e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.409892714606496, "grad_norm": 0.009115430043620344, "grad_rms": 0.0007759569355306323, "grad_to_weight": 0.0010838937371684249, "update_to_weight": 8.84483957667649e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3135810668666978, "grad_norm": 0.10138769420109707, "grad_rms": 0.0023470882893614022, "grad_to_weight": 0.07718419270684106, "update_to_weight": 0.00030985051674093576, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44550}
{"step": 44600, "lr": 1.2840667297704095e-05, "window_steps": 50, "elapsed_seconds": 2081.199184394, "steps_per_second": 1.7297719636807574, "eta_seconds": 3121.798776591, "images_seen": 45670400, "images_per_second": 1771.2864908090955, "effective_batch_size": 1024, "loss": 0.3645817041397095, "edge_loss": 0.18895308673381805, "diag_loss": 0.1756286323070526, "jvp_rms": 2.1816015243530273, "grad_norm": 0.21348786354064941, "grad_norm_postclip": 0.21348786354064941, "clip_fraction": 0.0, "ema_updates": 44600, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3735017641024156, "updates_nfe_2": 11, "loss_nfe_1": 0.407765719625685, "updates_nfe_1": 9, "loss_nfe_4": 0.3645570158958435, "updates_nfe_4": 15, "loss_nfe_5": 0.31654200951258343, "updates_nfe_5": 6, "loss_nfe_3": 0.3425629470083449, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.494208041787918, "grad_norm": 0.1579944388031078, "grad_rms": 0.0004600643735657971, "grad_to_weight": 0.00957878295234487, "update_to_weight": 4.7026870448381444e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25406349434273, "grad_norm": 0.11522895994232804, "grad_rms": 0.00033553547627438834, "grad_to_weight": 0.007089240175691069, "update_to_weight": 5.2287224733695245e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062268287159585, "grad_norm": 0.004679847882964721, "grad_rms": 8.37827169308247e-05, "grad_to_weight": 0.000924456709423151, "update_to_weight": 2.6627533098766893e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.852617381668619, "grad_norm": 0.05424085512740454, "grad_rms": 0.0001735709585786637, "grad_to_weight": 0.004220218615141145, "update_to_weight": 6.545970652078194e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.050583365709785, "grad_norm": 0.02428606574451043, "grad_rms": 0.0005525271272706153, "grad_to_weight": 0.004808566453807536, "update_to_weight": 2.513304247355893e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63971259969326, "grad_norm": 0.1335403260271837, "grad_rms": 0.0003217848523381678, "grad_to_weight": 0.008025398589495192, "update_to_weight": 7.276960807975972e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.565882947300484, "grad_norm": 0.0199233096379869, "grad_rms": 2.7487200035286596e-05, "grad_to_weight": 0.0008110154103040798, "update_to_weight": 8.973021196011565e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.827738926244454, "grad_norm": 0.009391468191508828, "grad_rms": 1.1401023597436642e-05, "grad_to_weight": 0.0003046434321368158, "update_to_weight": 7.157479601752111e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61384176875181, "grad_norm": 0.004363408830382827, "grad_rms": 3.917041870469859e-06, "grad_to_weight": 0.0001260596514982039, "update_to_weight": 8.54798870523894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.05693334038517, "grad_norm": 0.0035371881959491143, "grad_rms": 2.845179117678354e-06, "grad_to_weight": 9.545280402612908e-05, "update_to_weight": 8.518772027768894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.665868935272304, "grad_norm": 0.0014730499304058432, "grad_rms": 2.661346593779258e-06, "grad_to_weight": 5.972017179980127e-05, "update_to_weight": 5.4925574089637706e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31607892174297, "grad_norm": 0.003410263615392277, "grad_rms": 2.7430858316796143e-06, "grad_to_weight": 9.138858406168762e-05, "update_to_weight": 8.911200331445251e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.44251057798312, "grad_norm": 0.008699980184159573, "grad_rms": 3.862489764063368e-06, "grad_to_weight": 0.00013713171349777548, "update_to_weight": 9.99015250701482e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76604820537664, "grad_norm": 0.025837601006061413, "grad_rms": 1.9279268777230398e-05, "grad_to_weight": 0.0005089543488107263, "update_to_weight": 7.799667020557424e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.082522906849626, "grad_norm": 0.06237931445949278, "grad_rms": 5.3568893336507754e-05, "grad_to_weight": 0.001556272158927224, "update_to_weight": 9.16696095166214e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20008531922765, "grad_norm": 0.14935640706181244, "grad_rms": 0.00023282547350117748, "grad_to_weight": 0.005296310467542338, "update_to_weight": 7.52541102815857e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.409533458564908, "grad_norm": 0.003501403011917273, "grad_rms": 0.0002980592180712931, "grad_to_weight": 0.00041636114882820016, "update_to_weight": 4.490622921942196e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3132456069457428, "grad_norm": 0.051552367027521764, "grad_rms": 0.001193418569113173, "grad_to_weight": 0.03925569349317585, "update_to_weight": 0.0001434657428749484, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44600}
{"step": 44650, "lr": 1.2788863174884844e-05, "window_steps": 50, "elapsed_seconds": 2108.662461399, "steps_per_second": 1.7309550802068119, "eta_seconds": 3090.77922424237, "images_seen": 45721600, "images_per_second": 1772.4980021317754, "effective_batch_size": 1024, "loss": 0.3430500030517578, "edge_loss": 0.1758953183889389, "diag_loss": 0.16715466976165771, "jvp_rms": 1.693945288658142, "grad_norm": 0.14858423173427582, "grad_norm_postclip": 0.14858423173427582, "clip_fraction": 0.0, "ema_updates": 44650, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.35556461865251715, "updates_nfe_2": 11, "loss_nfe_3": 0.3179364750782649, "updates_nfe_3": 12, "loss_nfe_4": 0.38005094907500525, "updates_nfe_4": 11, "loss_nfe_5": 0.2903133794665337, "updates_nfe_5": 10, "loss_nfe_1": 0.3903928945461909, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.495776011473456, "grad_norm": 0.08022629300171691, "grad_rms": 0.00023361112905585994, "grad_to_weight": 0.004863444614301042, "update_to_weight": 5.023217790518781e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251561485974875, "grad_norm": 0.09147187285994582, "grad_rms": 0.00026635715918232325, "grad_to_weight": 0.005628497479389054, "update_to_weight": 5.635545303961639e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062121341813929, "grad_norm": 0.003392955565052284, "grad_rms": 6.074364867722023e-05, "grad_to_weight": 0.0006702635784381403, "update_to_weight": 3.0719387757079946e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.852490349294603, "grad_norm": 0.03709579834287887, "grad_rms": 0.00011870670664189413, "grad_to_weight": 0.0028862731917876786, "update_to_weight": 6.351674740762985e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0509141755516955, "grad_norm": 0.027784682663539464, "grad_rms": 0.0006321234182478162, "grad_to_weight": 0.005500921555552788, "update_to_weight": 2.3832748179845364e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63990700475676, "grad_norm": 0.13005216929512203, "grad_rms": 0.00031337963099154354, "grad_to_weight": 0.007815678853129691, "update_to_weight": 6.708511628888508e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.565665548242837, "grad_norm": 0.01731373046898823, "grad_rms": 2.3886893362874153e-05, "grad_to_weight": 0.0007047938691091831, "update_to_weight": 8.601523229355265e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828008188450767, "grad_norm": 0.007431675003720392, "grad_rms": 9.02188032351561e-06, "grad_to_weight": 0.0002410689318067767, "update_to_weight": 8.523099084276939e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61458411701273, "grad_norm": 0.002660651615283776, "grad_rms": 2.3884729084360043e-06, "grad_to_weight": 7.686504642926193e-05, "update_to_weight": 0.00010430545437588488, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.057526278831, "grad_norm": 0.0020926575507699764, "grad_rms": 1.6832538259404278e-06, "grad_to_weight": 5.647051384445488e-05, "update_to_weight": 0.00010633870724617722, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.6663526910268, "grad_norm": 0.0008339155295265289, "grad_rms": 1.506627988769939e-06, "grad_to_weight": 3.380781666313777e-05, "update_to_weight": 6.958935677457978e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31653198744911, "grad_norm": 0.002266282758992589, "grad_rms": 1.8229113135751737e-06, "grad_to_weight": 6.073133376260209e-05, "update_to_weight": 0.00010029308753484977, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.44363039960158, "grad_norm": 0.006527949946728335, "grad_rms": 2.8981835953448114e-06, "grad_to_weight": 0.0001028937011582069, "update_to_weight": 0.00010411239619512739, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76616362441252, "grad_norm": 0.021476081580404132, "grad_rms": 1.602482943262438e-05, "grad_to_weight": 0.000423039285365197, "update_to_weight": 7.627621461132344e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08288074853855, "grad_norm": 0.05798267765601038, "grad_rms": 4.979323517152336e-05, "grad_to_weight": 0.0014465696220729962, "update_to_weight": 8.433521920398107e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200170811345497, "grad_norm": 0.15293169276066487, "grad_rms": 0.00023839883725645852, "grad_to_weight": 0.00542307682402893, "update_to_weight": 6.81913263446951e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.409928548647686, "grad_norm": 0.008092369793525426, "grad_rms": 0.0006888682636053345, "grad_to_weight": 0.0009622400174644382, "update_to_weight": 3.6709122606143245e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.313664056378662, "grad_norm": 0.13921039410656388, "grad_rms": 0.003222670052213952, "grad_to_weight": 0.10597107641836603, "update_to_weight": 0.00010300160155535495, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44650}
{"step": 44700, "lr": 1.2737520666295111e-05, "window_steps": 50, "elapsed_seconds": 2136.1354777770002, "steps_per_second": 1.732099877789801, "eta_seconds": 3059.8697384373245, "images_seen": 45772800, "images_per_second": 1773.6702748567561, "effective_batch_size": 1024, "loss": 0.38194289803504944, "edge_loss": 0.19757381081581116, "diag_loss": 0.18436914682388306, "jvp_rms": 2.408281087875366, "grad_norm": 0.28168541193008423, "grad_norm_postclip": 0.27764227986335754, "clip_fraction": 0.019999999552965164, "ema_updates": 44700, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3991994640001884, "updates_nfe_4": 13, "loss_nfe_5": 0.37577543904383975, "updates_nfe_5": 6, "loss_nfe_1": 0.39769282124259253, "updates_nfe_1": 11, "loss_nfe_2": 0.35089510679244995, "updates_nfe_2": 10, "loss_nfe_3": 0.37693297266960146, "updates_nfe_3": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.495339400079896, "grad_norm": 0.2065472162770393, "grad_rms": 0.0006014453191398413, "grad_to_weight": 0.012521549952227045, "update_to_weight": 0.00013095088385601004, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25353026556846, "grad_norm": 0.08668908189490168, "grad_rms": 0.00025243013905491666, "grad_to_weight": 0.005333554032784136, "update_to_weight": 0.00010111450952328235, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0628010224692055, "grad_norm": 0.00596657798048727, "grad_rms": 0.00010681888097357739, "grad_to_weight": 0.00117851322894323, "update_to_weight": 5.675010594564018e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.853381923831673, "grad_norm": 0.06980265437763897, "grad_rms": 0.00022336877992066597, "grad_to_weight": 0.00543068390803954, "update_to_weight": 0.00016005637506911674, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.050784459891188, "grad_norm": 0.030322955380148807, "grad_rms": 0.0006898711220995424, "grad_to_weight": 0.006003613027035026, "update_to_weight": 4.688831158272095e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63968325533544, "grad_norm": 0.1535302565301998, "grad_rms": 0.00036995349941675796, "grad_to_weight": 0.009226753548987839, "update_to_weight": 0.00012677583888212368, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.565741327561362, "grad_norm": 0.023221631135765065, "grad_rms": 3.2037730265326994e-05, "grad_to_weight": 0.0009452851768699411, "update_to_weight": 0.00011362384677159748, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828283387834375, "grad_norm": 0.011531488487347743, "grad_rms": 1.3998958381894766e-05, "grad_to_weight": 0.0003740554847727383, "update_to_weight": 8.604830202704733e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61526649911405, "grad_norm": 0.002988895945870107, "grad_rms": 2.6831385784732087e-06, "grad_to_weight": 8.634617751524766e-05, "update_to_weight": 8.343417462414375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.058470019623215, "grad_norm": 0.0023550673952121605, "grad_rms": 1.8943262847186057e-06, "grad_to_weight": 6.355004386217522e-05, "update_to_weight": 8.767739433669647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66630443960001, "grad_norm": 0.0009197503410936722, "grad_rms": 1.6617050018953284e-06, "grad_to_weight": 3.728772355607019e-05, "update_to_weight": 5.669328415070154e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31685583609889, "grad_norm": 0.0024789804017037644, "grad_rms": 1.9939971755358965e-06, "grad_to_weight": 6.643058066284605e-05, "update_to_weight": 8.702003556941871e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.44401906159082, "grad_norm": 0.007510075704593213, "grad_rms": 3.33421340305437e-06, "grad_to_weight": 0.00011837326537120082, "update_to_weight": 9.524866613916749e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76638484349782, "grad_norm": 0.026458348642023167, "grad_rms": 1.974245266623501e-05, "grad_to_weight": 0.0005211785066750122, "update_to_weight": 9.180614327975606e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.083351267429286, "grad_norm": 0.07865722979726092, "grad_rms": 6.754772458890711e-05, "grad_to_weight": 0.001962341653333158, "update_to_weight": 0.00013578384888322665, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200046360708235, "grad_norm": 0.1770795034013103, "grad_rms": 0.000276041852089416, "grad_to_weight": 0.0062794046909099835, "update_to_weight": 0.00012136418493435877, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410255130340689, "grad_norm": 0.007362445771567841, "grad_rms": 0.0006267330045404258, "grad_to_weight": 0.0008754128926490246, "update_to_weight": 8.46933834039912e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3140970685950115, "grad_norm": 0.12036344824233312, "grad_rms": 0.002786370102040258, "grad_to_weight": 0.09159403145996033, "update_to_weight": 0.00010578372458026837, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44700}
{"step": 44750, "lr": 1.2686640333408315e-05, "window_steps": 50, "elapsed_seconds": 2163.596907877, "steps_per_second": 1.7332248841488855, "eta_seconds": 3029.0356710277997, "images_seen": 45824000, "images_per_second": 1774.8222813684588, "effective_batch_size": 1024, "loss": 0.36861899495124817, "edge_loss": 0.19150607287883759, "diag_loss": 0.1771128922700882, "jvp_rms": 2.0418357849121094, "grad_norm": 0.2512607276439667, "grad_norm_postclip": 0.24913345277309418, "clip_fraction": 0.019999999552965164, "ema_updates": 44750, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3553102132346895, "updates_nfe_3": 9, "loss_nfe_1": 0.3987983887394269, "updates_nfe_1": 12, "loss_nfe_4": 0.3899799548089504, "updates_nfe_4": 8, "loss_nfe_2": 0.3451028347015381, "updates_nfe_2": 10, "loss_nfe_5": 0.3524281518025832, "updates_nfe_5": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.496549303294177, "grad_norm": 0.01239703953888073, "grad_rms": 3.609896824680764e-05, "grad_to_weight": 0.0007514928916925178, "update_to_weight": 3.663274887878368e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25503884467204, "grad_norm": 0.023257086440344472, "grad_rms": 6.772236406040013e-05, "grad_to_weight": 0.0014307616648955295, "update_to_weight": 3.878859199578111e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062907261420472, "grad_norm": 0.0033370587885707442, "grad_rms": 5.9742935850988025e-05, "grad_to_weight": 0.0006591190824290319, "update_to_weight": 2.500912420815306e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854615898670165, "grad_norm": 0.012812536322437531, "grad_rms": 4.1000168712049635e-05, "grad_to_weight": 0.0009967265006932656, "update_to_weight": 5.2103568884034056e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.05052331913428, "grad_norm": 0.04949085380952099, "grad_rms": 0.0011259559110649955, "grad_to_weight": 0.009799153608898556, "update_to_weight": 2.7004693344317e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64022520059637, "grad_norm": 0.09238650796446936, "grad_rms": 0.00022261873778362754, "grad_to_weight": 0.005551998656914711, "update_to_weight": 7.050546671902404e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.565981086816226, "grad_norm": 0.030455295904158, "grad_rms": 4.201765800273038e-05, "grad_to_weight": 0.0012397345661273988, "update_to_weight": 8.692617722615996e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82799630928537, "grad_norm": 0.015313645937004646, "grad_rms": 1.8590409415264176e-05, "grad_to_weight": 0.0004967447700255558, "update_to_weight": 7.706921753408729e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61597884230544, "grad_norm": 0.01025826710544709, "grad_rms": 9.208869334156455e-06, "grad_to_weight": 0.00029634485138146927, "update_to_weight": 9.044747515011486e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.059260569327506, "grad_norm": 0.008336681593827217, "grad_rms": 6.705708338802767e-06, "grad_to_weight": 0.00022495542182315317, "update_to_weight": 9.438713660605137e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666045859832803, "grad_norm": 0.003369447429179104, "grad_rms": 6.087551584957813e-06, "grad_to_weight": 0.00013660265809632867, "update_to_weight": 6.530899708435725e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.317519880601154, "grad_norm": 0.008301319865675541, "grad_rms": 6.677264714961268e-06, "grad_to_weight": 0.00022245100671845116, "update_to_weight": 9.328770173259218e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.444929137032894, "grad_norm": 0.020296766361319517, "grad_rms": 9.011061020221791e-06, "grad_to_weight": 0.00031991156168653147, "update_to_weight": 9.97025695798157e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.766594038919926, "grad_norm": 0.031019929449345156, "grad_rms": 2.314617201358405e-05, "grad_to_weight": 0.0006110303445916403, "update_to_weight": 7.698046668627263e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08375325977627, "grad_norm": 0.03020412736334299, "grad_rms": 2.593811251484491e-05, "grad_to_weight": 0.0007535254288088982, "update_to_weight": 9.182865934773645e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200218427096225, "grad_norm": 0.0633459765346885, "grad_rms": 9.874740073908922e-05, "grad_to_weight": 0.002246293825647195, "update_to_weight": 6.974082169001298e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41029413786174, "grad_norm": 0.007627227475437491, "grad_rms": 0.0006492727200048727, "grad_to_weight": 0.0009068918815931759, "update_to_weight": 5.617999755659717e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3140991096976027, "grad_norm": 0.08044901010829379, "grad_rms": 0.0018623653590679084, "grad_to_weight": 0.06121989545126967, "update_to_weight": 0.00017891300751621844, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44750}
{"step": 44800, "lr": 1.2636222732643638e-05, "window_steps": 50, "elapsed_seconds": 2191.080426869, "steps_per_second": 1.7343042059985478, "eta_seconds": 2998.320584136526, "images_seen": 45875200, "images_per_second": 1775.927506942513, "effective_batch_size": 1024, "loss": 0.37544330954551697, "edge_loss": 0.19241949915885925, "diag_loss": 0.18302379548549652, "jvp_rms": 2.555117130279541, "grad_norm": 0.17152836918830872, "grad_norm_postclip": 0.17152836918830872, "clip_fraction": 0.0, "ema_updates": 44800, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.39163410663604736, "updates_nfe_2": 12, "loss_nfe_4": 0.38539310097694396, "updates_nfe_4": 15, "loss_nfe_5": 0.2911337961753209, "updates_nfe_5": 6, "loss_nfe_3": 0.3760807920585979, "updates_nfe_3": 11, "loss_nfe_1": 0.4013281712929408, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.496304184810988, "grad_norm": 0.06667514662653187, "grad_rms": 0.0001941515144300109, "grad_to_weight": 0.004041823300513771, "update_to_weight": 3.5885948971998726e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25176898424253, "grad_norm": 0.07794346580026848, "grad_rms": 0.00022696375922215238, "grad_to_weight": 0.004795998877158621, "update_to_weight": 3.842944249812412e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063274936469023, "grad_norm": 0.003250733869799751, "grad_rms": 5.8197471892684175e-05, "grad_to_weight": 0.0006420219937862422, "update_to_weight": 2.6478646663432068e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.853278642065842, "grad_norm": 0.030139876504313344, "grad_rms": 9.644772826697389e-05, "grad_to_weight": 0.0023449173820656484, "update_to_weight": 4.6803682382774546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.050493484406762, "grad_norm": 0.01643034706546797, "grad_rms": 0.00037380333890408825, "grad_to_weight": 0.003253216169112215, "update_to_weight": 2.0276680215938985e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639796047427986, "grad_norm": 0.09517038236476848, "grad_rms": 0.00022932688834367563, "grad_to_weight": 0.005719444042072797, "update_to_weight": 6.080375611946798e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.565880462750695, "grad_norm": 0.015983636168611905, "grad_rms": 2.2051828367923205e-05, "grad_to_weight": 0.0006506437329957675, "update_to_weight": 7.118873861866769e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82824973048766, "grad_norm": 0.007072854827694785, "grad_rms": 8.58628098902571e-06, "grad_to_weight": 0.00022942771287790857, "update_to_weight": 6.337869398836587e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.616359693047414, "grad_norm": 0.0035517730348741733, "grad_rms": 3.188434600749369e-06, "grad_to_weight": 0.00010260388632336565, "update_to_weight": 7.984349999878439e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.060148270666026, "grad_norm": 0.002945757171193291, "grad_rms": 2.36945458509355e-06, "grad_to_weight": 7.948584419250494e-05, "update_to_weight": 8.521865159387004e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66629949073039, "grad_norm": 0.0010300366392220622, "grad_rms": 1.8609583047237323e-06, "grad_to_weight": 4.175886373264667e-05, "update_to_weight": 5.6750660836916476e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31852410279562, "grad_norm": 0.0031001113754645735, "grad_rms": 2.4936112130110174e-06, "grad_to_weight": 8.307165007182951e-05, "update_to_weight": 8.824250965053606e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.44640485529893, "grad_norm": 0.008301431616381057, "grad_rms": 3.6855480089166766e-06, "grad_to_weight": 0.00013084163925937146, "update_to_weight": 9.669715772876618e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.766916246274505, "grad_norm": 0.021183408571726813, "grad_rms": 1.5806445318835834e-05, "grad_to_weight": 0.0004172679795826941, "update_to_weight": 7.150084941055351e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.083884210958146, "grad_norm": 0.04391548150899686, "grad_rms": 3.771288230979959e-05, "grad_to_weight": 0.00109558947126664, "update_to_weight": 7.739789437755169e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200238988418278, "grad_norm": 0.09942448510804491, "grad_rms": 0.00015498868296497716, "grad_to_weight": 0.003525661082123387, "update_to_weight": 5.665882657461657e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410448805437598, "grad_norm": 0.003053280927131618, "grad_rms": 0.00025991253294619163, "grad_to_weight": 0.00036303424439818044, "update_to_weight": 3.2905886999021588e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3142463782225933, "grad_norm": 0.052786084272037584, "grad_rms": 0.0012219786751477668, "grad_to_weight": 0.04016452709835601, "update_to_weight": 0.00011836808185731313, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44800}
{"step": 44850, "lr": 1.2586268415359876e-05, "window_steps": 50, "elapsed_seconds": 2218.557769097, "steps_per_second": 1.7353616180871556, "eta_seconds": 2967.681171649234, "images_seen": 45926400, "images_per_second": 1777.0102969212473, "effective_batch_size": 1024, "loss": 0.3489411771297455, "edge_loss": 0.18235638737678528, "diag_loss": 0.1665848195552826, "jvp_rms": 1.6544530391693115, "grad_norm": 0.19229769706726074, "grad_norm_postclip": 0.19229769706726074, "clip_fraction": 0.0, "ema_updates": 44850, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3377015675817217, "updates_nfe_4": 7, "loss_nfe_2": 0.33226993680000305, "updates_nfe_2": 11, "loss_nfe_1": 0.4044209374831273, "updates_nfe_1": 13, "loss_nfe_5": 0.29482054048114353, "updates_nfe_5": 9, "loss_nfe_3": 0.3517321854829788, "updates_nfe_3": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49501008384098, "grad_norm": 0.055894716793431894, "grad_rms": 0.00016275995574283503, "grad_to_weight": 0.0033885833660803927, "update_to_weight": 4.200844172735822e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251217840482852, "grad_norm": 0.03265939112208345, "grad_rms": 9.510095691625143e-05, "grad_to_weight": 0.0020096580725615996, "update_to_weight": 6.298088977216375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062865255860462, "grad_norm": 0.0032129832631418893, "grad_rms": 5.75216276193915e-05, "grad_to_weight": 0.0006346175734032694, "update_to_weight": 2.7280407952186453e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.852515280981056, "grad_norm": 0.019824770657838525, "grad_rms": 6.34393473074998e-05, "grad_to_weight": 0.0015424817807588916, "update_to_weight": 7.041729378120955e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.050521997346285, "grad_norm": 0.030322599293996822, "grad_rms": 0.0006898630208590744, "grad_to_weight": 0.006003854514430258, "update_to_weight": 2.835099353955056e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640186687302556, "grad_norm": 0.057005328826259964, "grad_rms": 0.00013736263692446535, "grad_to_weight": 0.003425762576916183, "update_to_weight": 6.242104825625101e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566491654426997, "grad_norm": 0.029823803842577615, "grad_rms": 4.114641979317809e-05, "grad_to_weight": 0.001214003377531656, "update_to_weight": 8.899518893899087e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828595211505238, "grad_norm": 0.01312460384399009, "grad_rms": 1.593296331106442e-05, "grad_to_weight": 0.0004257282485285604, "update_to_weight": 7.74716208241874e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61658890675196, "grad_norm": 0.009147416160487547, "grad_rms": 8.211655955258889e-06, "grad_to_weight": 0.00026424949567180914, "update_to_weight": 9.619369651271619e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.060615993430275, "grad_norm": 0.007617794277861816, "grad_rms": 6.127462832473473e-06, "grad_to_weight": 0.00020554958609463533, "update_to_weight": 9.611209497277452e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66690448268312, "grad_norm": 0.003239442849574185, "grad_rms": 5.852673433195548e-06, "grad_to_weight": 0.0001313274980185847, "update_to_weight": 6.240923959987764e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3193238630304, "grad_norm": 0.007617421053017139, "grad_rms": 6.127162624659964e-06, "grad_to_weight": 0.00020411465869463882, "update_to_weight": 9.927312085717344e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.447790113702936, "grad_norm": 0.019863682873406216, "grad_rms": 8.818786957104304e-06, "grad_to_weight": 0.0003130713116691549, "update_to_weight": 0.00011313903205244268, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76710379886763, "grad_norm": 0.029404933647626738, "grad_rms": 2.194110896891037e-05, "grad_to_weight": 0.0005792123530253979, "update_to_weight": 8.629919708054623e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.084612048102386, "grad_norm": 0.03394977010797798, "grad_rms": 2.9154722674841328e-05, "grad_to_weight": 0.0008469526926501754, "update_to_weight": 9.798761176011716e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200048525071836, "grad_norm": 0.057364394164467626, "grad_rms": 8.942296146641306e-05, "grad_to_weight": 0.00203419487429132, "update_to_weight": 7.978445841845941e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410529539875196, "grad_norm": 0.005753375433730833, "grad_rms": 0.000489759841187032, "grad_to_weight": 0.0006840681560481396, "update_to_weight": 9.046801392966506e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3145327040912398, "grad_norm": 0.0472292330438809, "grad_rms": 0.001093339587868978, "grad_to_weight": 0.0359285340690944, "update_to_weight": 0.00020494506789704793, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44850}
{"step": 44900, "lr": 1.253677792784943e-05, "window_steps": 50, "elapsed_seconds": 2246.024379615, "steps_per_second": 1.7364014546754452, "eta_seconds": 2937.108804111923, "images_seen": 45977600, "images_per_second": 1778.0750895876558, "effective_batch_size": 1024, "loss": 0.36594679951667786, "edge_loss": 0.18911096453666687, "diag_loss": 0.1768358200788498, "jvp_rms": 2.1468358039855957, "grad_norm": 0.21371613442897797, "grad_norm_postclip": 0.21371613442897797, "clip_fraction": 0.0, "ema_updates": 44900, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.37685955233044094, "updates_nfe_2": 9, "loss_nfe_3": 0.333163166588003, "updates_nfe_3": 11, "loss_nfe_4": 0.3930546858093955, "updates_nfe_4": 11, "loss_nfe_1": 0.39977973997592925, "updates_nfe_1": 10, "loss_nfe_5": 0.3243790600034926, "updates_nfe_5": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.496369858422167, "grad_norm": 0.08676684695372082, "grad_rms": 0.00025265658331043584, "grad_to_weight": 0.005259753976079913, "update_to_weight": 5.95871747471408e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25259895082168, "grad_norm": 0.025760977986019867, "grad_rms": 7.501345167186601e-05, "grad_to_weight": 0.001585037449331602, "update_to_weight": 6.603804090490496e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062643166860661, "grad_norm": 0.0035587614199023457, "grad_rms": 6.371204964874404e-05, "grad_to_weight": 0.0007029453395407145, "update_to_weight": 3.7030171880365826e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.853855584747896, "grad_norm": 0.03225381558590214, "grad_rms": 0.00010321234198676913, "grad_to_weight": 0.00250927166352902, "update_to_weight": 8.685368432929145e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.050140174959816, "grad_norm": 0.039677628449644356, "grad_rms": 0.0009026973036646849, "grad_to_weight": 0.007856738045882076, "update_to_weight": 2.460908790236475e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640072063876875, "grad_norm": 0.11776288556200924, "grad_rms": 0.00028376681313308964, "grad_to_weight": 0.007077065839014903, "update_to_weight": 9.081404485476892e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566404697189487, "grad_norm": 0.05957793945203071, "grad_rms": 8.219672178791834e-05, "grad_to_weight": 0.0024251794345326694, "update_to_weight": 0.00011184894173481802, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828694202487213, "grad_norm": 0.04703030092606153, "grad_rms": 5.709368968926139e-05, "grad_to_weight": 0.0015255365866993873, "update_to_weight": 0.00010317151094653303, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.616613591214225, "grad_norm": 0.02858851918791366, "grad_rms": 2.5663977643819346e-05, "grad_to_weight": 0.0008258612331498968, "update_to_weight": 0.00011731078218220841, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.060726335616216, "grad_norm": 0.020039514273715295, "grad_rms": 1.6119020075123048e-05, "grad_to_weight": 0.0005407210342355557, "update_to_weight": 0.00011809418693999964, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66664467200596, "grad_norm": 0.0073050446665239195, "grad_rms": 1.3197961141278128e-05, "grad_to_weight": 0.0002961507235239974, "update_to_weight": 7.957222461774868e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.31970819978618, "grad_norm": 0.016508628554750004, "grad_rms": 1.3278910429271092e-05, "grad_to_weight": 0.00044235684980100114, "update_to_weight": 0.00011495678589437531, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.44905028778504, "grad_norm": 0.028591270777374632, "grad_rms": 1.2693533592208016e-05, "grad_to_weight": 0.0004506177893552949, "update_to_weight": 0.0001222275517322641, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76751977842846, "grad_norm": 0.034992706859321716, "grad_rms": 2.6110543336644542e-05, "grad_to_weight": 0.0006892735160599753, "update_to_weight": 8.989757211152846e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08466229578575, "grad_norm": 0.04702449840420607, "grad_rms": 4.038278331598992e-05, "grad_to_weight": 0.001173129464262692, "update_to_weight": 0.00010355610224724489, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200326644412517, "grad_norm": 0.11769792736173383, "grad_rms": 0.00018347438993200922, "grad_to_weight": 0.004173637023635468, "update_to_weight": 8.325029317859344e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410473297989324, "grad_norm": 0.005258182087264524, "grad_rms": 0.0004476061841702594, "grad_to_weight": 0.000625194552192632, "update_to_weight": 3.7884182202248547e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3144452348446791, "grad_norm": 0.08978694695742256, "grad_rms": 0.0020785352049065495, "grad_to_weight": 0.0683078644718372, "update_to_weight": 0.00010852726687107308, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44900}
{"step": 44950, "lr": 1.248775181133237e-05, "window_steps": 50, "elapsed_seconds": 2273.513596031, "steps_per_second": 1.7373988908162838, "eta_seconds": 2906.643964545962, "images_seen": 46028800, "images_per_second": 1779.0964641958747, "effective_batch_size": 1024, "loss": 0.34026065468788147, "edge_loss": 0.17515771090984344, "diag_loss": 0.16510285437107086, "jvp_rms": 1.60476553440094, "grad_norm": 0.2855491042137146, "grad_norm_postclip": 0.2652169167995453, "clip_fraction": 0.03999999910593033, "ema_updates": 44950, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39569009840488434, "updates_nfe_1": 8, "loss_nfe_3": 0.33507657547791797, "updates_nfe_3": 12, "loss_nfe_4": 0.2879256470636888, "updates_nfe_4": 11, "loss_nfe_5": 0.33333532388011616, "updates_nfe_5": 12, "loss_nfe_2": 0.3799120954104832, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.496552078199553, "grad_norm": 0.02408167836858103, "grad_rms": 7.012349521277143e-05, "grad_to_weight": 0.0014598007058944963, "update_to_weight": 4.413159551856684e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25442116088391, "grad_norm": 0.006566721854576435, "grad_rms": 1.9121652630896972e-05, "grad_to_weight": 0.0004039960444964463, "update_to_weight": 4.1604193151017204e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062397331797961, "grad_norm": 0.002738681579900801, "grad_rms": 4.9030265365619435e-05, "grad_to_weight": 0.000540985110492726, "update_to_weight": 2.7459748766233663e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.853961829599342, "grad_norm": 0.007316230600052187, "grad_rms": 2.341196788750507e-05, "grad_to_weight": 0.0005691809807000363, "update_to_weight": 6.615965697882098e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.050019315172781, "grad_norm": 0.023314634529255536, "grad_rms": 0.0005304263018692382, "grad_to_weight": 0.00461674165467184, "update_to_weight": 2.368001116310108e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640125249244576, "grad_norm": 0.061456297388473614, "grad_rms": 0.00014808789351296726, "grad_to_weight": 0.0036932593035177777, "update_to_weight": 6.557476600157927e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56607798367497, "grad_norm": 0.02529618807044662, "grad_rms": 3.4899893353906325e-05, "grad_to_weight": 0.001029720254379101, "update_to_weight": 7.70256963434089e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828765475797418, "grad_norm": 0.014848483579438276, "grad_rms": 1.8025713149769884e-05, "grad_to_weight": 0.0004816437943678056, "update_to_weight": 5.945130438829586e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61748811509058, "grad_norm": 0.010647050027472129, "grad_rms": 9.557880633187558e-06, "grad_to_weight": 0.0003075627553355273, "update_to_weight": 6.842436228456452e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06121710732882, "grad_norm": 0.008752956085953703, "grad_rms": 7.0405436448735105e-06, "grad_to_weight": 0.00023617562425446665, "update_to_weight": 7.080597538252529e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666984899719072, "grad_norm": 0.00337823912547857, "grad_rms": 6.1034354667714314e-06, "grad_to_weight": 0.00013695387333362516, "update_to_weight": 4.5063092573614246e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.321001831869076, "grad_norm": 0.008932364332195002, "grad_rms": 7.184852787465815e-06, "grad_to_weight": 0.0002393388144411358, "update_to_weight": 7.602837168675657e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45100302718823, "grad_norm": 0.02070549223744535, "grad_rms": 9.192521147650389e-06, "grad_to_weight": 0.0003263225362816285, "update_to_weight": 8.725483138917558e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.767599126743725, "grad_norm": 0.026884689472976064, "grad_rms": 2.0060575833657573e-05, "grad_to_weight": 0.0005295639332058456, "update_to_weight": 6.79063280836504e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08508559195131, "grad_norm": 0.025290165519390832, "grad_rms": 2.171819602234676e-05, "grad_to_weight": 0.0006309120997478635, "update_to_weight": 8.085631437902525e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200435943480443, "grad_norm": 0.04191334525389665, "grad_rms": 6.533696576349871e-05, "grad_to_weight": 0.0014862658626235332, "update_to_weight": 6.906723619456088e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410570360363266, "grad_norm": 0.005751330461076984, "grad_rms": 0.0004895857615543508, "grad_to_weight": 0.0006838216927809608, "update_to_weight": 4.761315765985161e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3146074269661674, "grad_norm": 0.08058992100502488, "grad_rms": 0.0018656273951381337, "grad_to_weight": 0.06130341222170726, "update_to_weight": 0.00011392386551343138, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 44950}
{"step": 45000, "lr": 1.2439190601950434e-05, "window_steps": 50, "elapsed_seconds": 2301.024512436, "steps_per_second": 1.738356101111398, "eta_seconds": 2876.2806405449996, "images_seen": 46080000, "images_per_second": 1780.0766475380715, "effective_batch_size": 1024, "loss": 0.3361496031284332, "edge_loss": 0.17407377064228058, "diag_loss": 0.16207581758499146, "jvp_rms": 1.5733202695846558, "grad_norm": 0.182581827044487, "grad_norm_postclip": 0.182581827044487, "clip_fraction": 0.0, "ema_updates": 45000, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3178708576239072, "updates_nfe_5": 13, "loss_nfe_1": 0.39370010495185853, "updates_nfe_1": 10, "loss_nfe_2": 0.36512787852968487, "updates_nfe_2": 7, "loss_nfe_3": 0.31815891795688206, "updates_nfe_3": 9, "loss_nfe_4": 0.30171205509792676, "updates_nfe_4": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.498510120057528, "grad_norm": 0.012333655460575039, "grad_rms": 3.591440000187028e-05, "grad_to_weight": 0.0007475617719918114, "update_to_weight": 6.054606717789965e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.251992440952623, "grad_norm": 0.00966916802342172, "grad_rms": 2.815567284074842e-05, "grad_to_weight": 0.0005949527763166344, "update_to_weight": 9.736419165062931e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062393752501982, "grad_norm": 0.0007599230084486812, "grad_rms": 1.3604804236872357e-05, "grad_to_weight": 0.00015011139899441943, "update_to_weight": 3.412245941339771e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854061544638697, "grad_norm": 0.007665439099224211, "grad_rms": 2.4529436515216306e-05, "grad_to_weight": 0.0005963437371607569, "update_to_weight": 7.58617391476943e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0499180928754654, "grad_norm": 0.018577945400060508, "grad_rms": 0.00042266289280741207, "grad_to_weight": 0.003678860737616849, "update_to_weight": 1.9419408936065435e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64048745541747, "grad_norm": 0.07596542428754491, "grad_rms": 0.0001830497465776588, "grad_to_weight": 0.004565096094153998, "update_to_weight": 6.529756139441512e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566432026640157, "grad_norm": 0.013721362132913553, "grad_rms": 1.893068132540006e-05, "grad_to_weight": 0.0005585411067441104, "update_to_weight": 9.942331102400385e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828619959280534, "grad_norm": 0.006566901625851995, "grad_rms": 7.972065588858176e-06, "grad_to_weight": 0.00021301315577946004, "update_to_weight": 9.116016129398156e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.616856908543305, "grad_norm": 0.003688622526991926, "grad_rms": 3.3112846960338723e-06, "grad_to_weight": 0.00010655567421205096, "update_to_weight": 0.00011250295535780627, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06090419994077, "grad_norm": 0.0031675091869591008, "grad_rms": 2.547823438999179e-06, "grad_to_weight": 8.546767153522829e-05, "update_to_weight": 0.00010773269710218021, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666250001979552, "grad_norm": 0.0010668607778810295, "grad_rms": 1.9274881581699743e-06, "grad_to_weight": 4.325184321878722e-05, "update_to_weight": 7.280571413098203e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.321369797405445, "grad_norm": 0.003186210334280857, "grad_rms": 2.562865927803541e-06, "grad_to_weight": 8.5372277372905e-05, "update_to_weight": 0.00010569143558248016, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.451587875373534, "grad_norm": 0.008584917081414511, "grad_rms": 3.811405733161394e-06, "grad_to_weight": 0.0001352986957280929, "update_to_weight": 0.00011543422669937828, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76762317166325, "grad_norm": 0.020084192360247437, "grad_rms": 1.4986242050721731e-05, "grad_to_weight": 0.0003956102552277402, "update_to_weight": 7.933160353679642e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08533987203115, "grad_norm": 0.030192573860835615, "grad_rms": 2.5928190822873087e-05, "grad_to_weight": 0.0007532073809832397, "update_to_weight": 8.208681359096637e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200603678849898, "grad_norm": 0.06156059308777861, "grad_rms": 9.596424094979456e-05, "grad_to_weight": 0.002182953024298139, "update_to_weight": 5.535411633915595e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410573081722093, "grad_norm": 0.006546257810581789, "grad_rms": 0.0005572544713288215, "grad_to_weight": 0.0007783367134408659, "update_to_weight": 2.5839049461208393e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3145658946303826, "grad_norm": 0.08024733680598307, "grad_rms": 0.0018576966953818452, "grad_to_weight": 0.061044742704622096, "update_to_weight": 8.440423355000197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45000}
{"step": 45050, "lr": 1.239109483076124e-05, "window_steps": 50, "elapsed_seconds": 2351.833748903, "steps_per_second": 1.7220604993397601, "eta_seconds": 2874.463470881444, "images_seen": 46131200, "images_per_second": 1763.3899513239144, "effective_batch_size": 1024, "loss": 0.36284229159355164, "edge_loss": 0.18799163401126862, "diag_loss": 0.1748506873846054, "jvp_rms": 1.9524608850479126, "grad_norm": 0.2414599359035492, "grad_norm_postclip": 0.22997204959392548, "clip_fraction": 0.019999999552965164, "ema_updates": 45050, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3103618941136769, "updates_nfe_3": 7, "loss_nfe_4": 0.3681761850329006, "updates_nfe_4": 17, "loss_nfe_1": 0.3990859157509274, "updates_nfe_1": 9, "loss_nfe_5": 0.39219771538461956, "updates_nfe_5": 7, "loss_nfe_2": 0.33734299540519713, "updates_nfe_2": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.493947161572635, "grad_norm": 0.007724807880406425, "grad_rms": 2.2493886021166767e-05, "grad_to_weight": 0.0004683419805298985, "update_to_weight": 5.025840796226229e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.254150799500128, "grad_norm": 0.01752910799865125, "grad_rms": 5.104305032290832e-05, "grad_to_weight": 0.0010784388686236584, "update_to_weight": 7.82116477146886e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062523735838454, "grad_norm": 0.0025971044272978428, "grad_rms": 4.6495627745541205e-05, "grad_to_weight": 0.0005130058766761932, "update_to_weight": 2.4939237891552912e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.853886449264516, "grad_norm": 0.008036603448901619, "grad_rms": 2.571716395447611e-05, "grad_to_weight": 0.0006252275123654499, "update_to_weight": 7.040244010907034e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.049702799899175, "grad_norm": 0.017639324326064715, "grad_rms": 0.00040130852396618545, "grad_to_weight": 0.003493141086722353, "update_to_weight": 2.3999657117870312e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639686006371186, "grad_norm": 0.03334234131942882, "grad_rms": 8.034322438225266e-05, "grad_to_weight": 0.002003784284550943, "update_to_weight": 6.190968775406014e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566089164057153, "grad_norm": 0.02020912700019132, "grad_rms": 2.7881528043595597e-05, "grad_to_weight": 0.000822643232515799, "update_to_weight": 7.740587502389279e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828545715895054, "grad_norm": 0.012606455396887106, "grad_rms": 1.5303943167255873e-05, "grad_to_weight": 0.0004089215077825509, "update_to_weight": 6.349592739289014e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61807523263678, "grad_norm": 0.009332614175873142, "grad_rms": 8.377908628064166e-06, "grad_to_weight": 0.0002695878991872622, "update_to_weight": 7.20986754011688e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.062252976621984, "grad_norm": 0.007617598235613611, "grad_rms": 6.127305143574868e-06, "grad_to_weight": 0.00020553521774347114, "update_to_weight": 7.691819518816179e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66660136996038, "grad_norm": 0.003129957508270395, "grad_rms": 5.6548672121482264e-06, "grad_to_weight": 0.00012689050515415302, "update_to_weight": 5.228728544493176e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.322810551993534, "grad_norm": 0.007642633871880031, "grad_rms": 6.147442853404553e-06, "grad_to_weight": 0.000204771124115459, "update_to_weight": 8.131033253823728e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45341164365199, "grad_norm": 0.018460640857079297, "grad_rms": 8.195884914582346e-06, "grad_to_weight": 0.0002909322033108702, "update_to_weight": 9.22758133315422e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76811368553592, "grad_norm": 0.02630680397625036, "grad_rms": 1.962937443027556e-05, "grad_to_weight": 0.0005181757222495602, "update_to_weight": 7.260229664307185e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.085930648264856, "grad_norm": 0.02596623154333602, "grad_rms": 2.2298774841446707e-05, "grad_to_weight": 0.0006477642186027177, "update_to_weight": 8.28604535062106e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200192454878522, "grad_norm": 0.053292436526612544, "grad_rms": 8.307535654098134e-05, "grad_to_weight": 0.0018897898165724455, "update_to_weight": 6.61320018836911e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410774459831874, "grad_norm": 0.006295202934566679, "grad_rms": 0.0005358832610501524, "grad_to_weight": 0.0007484688793679193, "update_to_weight": 4.7585222481982486e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3149632998519694, "grad_norm": 0.07178971774298004, "grad_rms": 0.0016619058864964867, "grad_to_weight": 0.05459446491857354, "update_to_weight": 7.526910831466346e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45050}
{"step": 45100, "lr": 1.2343465023732415e-05, "window_steps": 50, "elapsed_seconds": 2379.307585034, "steps_per_second": 1.7231904045484776, "eta_seconds": 2843.56272357722, "images_seen": 46182400, "images_per_second": 1764.546974257641, "effective_batch_size": 1024, "loss": 0.37690016627311707, "edge_loss": 0.19406352937221527, "diag_loss": 0.18283675611019135, "jvp_rms": 2.169179677963257, "grad_norm": 0.2274029552936554, "grad_norm_postclip": 0.2274029552936554, "clip_fraction": 0.0, "ema_updates": 45100, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.39855845485414776, "updates_nfe_5": 7, "loss_nfe_4": 0.3933319775387645, "updates_nfe_4": 16, "loss_nfe_1": 0.40115777935300556, "updates_nfe_1": 7, "loss_nfe_2": 0.3518690994152656, "updates_nfe_2": 13, "loss_nfe_3": 0.33991297440869467, "updates_nfe_3": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.495987837837834, "grad_norm": 0.0336877917686701, "grad_rms": 9.809555914928422e-05, "grad_to_weight": 0.0020421809290740623, "update_to_weight": 5.461666426161673e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25586395098608, "grad_norm": 0.013887357063350812, "grad_rms": 4.043862731014912e-05, "grad_to_weight": 0.0008542983076890481, "update_to_weight": 4.606850708667506e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062559151009567, "grad_norm": 0.0031306423241249355, "grad_rms": 5.604748834008114e-05, "grad_to_weight": 0.0006183912584015197, "update_to_weight": 2.6616090250259706e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.855321567449764, "grad_norm": 0.014443848345657765, "grad_rms": 4.622037386822126e-05, "grad_to_weight": 0.0011235695871062627, "update_to_weight": 6.933029315298256e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.049200225045315, "grad_norm": 0.023603235802591768, "grad_rms": 0.0005369922081860797, "grad_to_weight": 0.004674648409764724, "update_to_weight": 3.3009866499304575e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.63995835665847, "grad_norm": 0.11572740038068398, "grad_rms": 0.00027886201532410456, "grad_to_weight": 0.006954789062580539, "update_to_weight": 0.00010323542534016243, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566608425090322, "grad_norm": 0.01697978531062333, "grad_rms": 2.342616582635633e-05, "grad_to_weight": 0.0006911733608812508, "update_to_weight": 0.00010480296630554375, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828305166097966, "grad_norm": 0.008085206984004755, "grad_rms": 9.815252922662376e-06, "grad_to_weight": 0.00026226569837177086, "update_to_weight": 7.956489912323688e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.618008234741424, "grad_norm": 0.002546903633997234, "grad_rms": 2.28636109111594e-06, "grad_to_weight": 7.357163984498826e-05, "update_to_weight": 7.852520519907648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06237154804891, "grad_norm": 0.002213294207730967, "grad_rms": 1.7802893463023406e-06, "grad_to_weight": 5.9718094533200014e-05, "update_to_weight": 8.183358885196034e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.666288355770106, "grad_norm": 0.0008394004441493917, "grad_rms": 1.5165375366728427e-06, "grad_to_weight": 3.403026965558981e-05, "update_to_weight": 5.644494043800893e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32271897333446, "grad_norm": 0.002185230399881606, "grad_rms": 1.7577158908817375e-06, "grad_to_weight": 5.854960356566901e-05, "update_to_weight": 8.431765459273657e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.4539887745936, "grad_norm": 0.006482154845558339, "grad_rms": 2.877852157138157e-06, "grad_to_weight": 0.00010215519892035117, "update_to_weight": 9.493071962370664e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76812330340762, "grad_norm": 0.018840271918900506, "grad_rms": 1.4058064681600177e-05, "grad_to_weight": 0.000371104360236139, "update_to_weight": 8.504089268608819e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.086248872061006, "grad_norm": 0.03772558549481698, "grad_rms": 3.239724390914038e-05, "grad_to_weight": 0.000941110394619903, "update_to_weight": 0.00010851594035621613, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200756263061347, "grad_norm": 0.12279329965194015, "grad_rms": 0.0001914173532736561, "grad_to_weight": 0.00435425555635827, "update_to_weight": 9.535317228645548e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410437012702069, "grad_norm": 0.010344603033299657, "grad_rms": 0.0008805910890203094, "grad_to_weight": 0.0012299721188894783, "update_to_weight": 2.992152829539093e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3145739654271715, "grad_norm": 0.170948996558275, "grad_rms": 0.003957407169199321, "grad_to_weight": 0.13004136781510428, "update_to_weight": 0.00011931810453454667, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45100}
{"step": 45150, "lr": 1.2296301701735908e-05, "window_steps": 50, "elapsed_seconds": 2406.790702736, "steps_per_second": 1.724287864035019, "eta_seconds": 2812.7553995830363, "images_seen": 46233600, "images_per_second": 1765.6707727718594, "effective_batch_size": 1024, "loss": 0.35431697964668274, "edge_loss": 0.18314507603645325, "diag_loss": 0.1711718589067459, "jvp_rms": 1.7952734231948853, "grad_norm": 0.21674302220344543, "grad_norm_postclip": 0.21674302220344543, "clip_fraction": 0.0, "ema_updates": 45150, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39715055955780876, "updates_nfe_1": 9, "loss_nfe_2": 0.3509938581423326, "updates_nfe_2": 11, "loss_nfe_4": 0.2938152568207847, "updates_nfe_4": 9, "loss_nfe_3": 0.3820884845086506, "updates_nfe_3": 14, "loss_nfe_5": 0.3267118015459606, "updates_nfe_5": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.498806996998866, "grad_norm": 0.009914426814555724, "grad_rms": 2.8869842484688764e-05, "grad_to_weight": 0.0006009178006845682, "update_to_weight": 6.414204775223536e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25164880457326, "grad_norm": 0.006937551649011179, "grad_rms": 2.0201472771204328e-05, "grad_to_weight": 0.00042688294169013376, "update_to_weight": 0.00011660541774471988, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062669350888017, "grad_norm": 0.0007107992415807688, "grad_rms": 1.2725347733798417e-05, "grad_to_weight": 0.00014040009179270045, "update_to_weight": 2.7861940627616392e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854021777424068, "grad_norm": 0.006003649101853653, "grad_rms": 1.9211701716925624e-05, "grad_to_weight": 0.00046706386575430075, "update_to_weight": 8.650817896999597e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.049996087101091, "grad_norm": 0.01694375334509764, "grad_rms": 0.0003854837362064219, "grad_to_weight": 0.0033552012819130844, "update_to_weight": 2.5640209804191146e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640001455453664, "grad_norm": 0.06526419720287743, "grad_rms": 0.00015726358235505563, "grad_to_weight": 0.003922126892692517, "update_to_weight": 7.024178435630345e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56640718168626, "grad_norm": 0.012247959308382718, "grad_rms": 1.6897900682709233e-05, "grad_to_weight": 0.0004985653464831159, "update_to_weight": 8.688117746023375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828730829069983, "grad_norm": 0.0066518501226448075, "grad_rms": 8.075191084973606e-06, "grad_to_weight": 0.00021576788741404945, "update_to_weight": 7.47152263359035e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61921417701989, "grad_norm": 0.003735046314836687, "grad_rms": 3.352959434258616e-06, "grad_to_weight": 0.00010788940198752395, "update_to_weight": 8.522864120328259e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06311425766273, "grad_norm": 0.0031345513610294726, "grad_rms": 2.5213134223123524e-06, "grad_to_weight": 8.45733399313958e-05, "update_to_weight": 8.507060153555427e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.667127175371224, "grad_norm": 0.0010595308242878636, "grad_rms": 1.914245194285948e-06, "grad_to_weight": 4.295315043195412e-05, "update_to_weight": 5.7117823297216224e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3232030294159, "grad_norm": 0.0031869768104461245, "grad_rms": 2.563482451963074e-06, "grad_to_weight": 8.538862026215546e-05, "update_to_weight": 8.625684892632719e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.455306537234335, "grad_norm": 0.00883108019553553, "grad_rms": 3.920693626749194e-06, "grad_to_weight": 0.00013917008170709293, "update_to_weight": 9.492540748873879e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76859217244334, "grad_norm": 0.01813491272365954, "grad_rms": 1.353174610015172e-05, "grad_to_weight": 0.0003572073194793647, "update_to_weight": 7.212026068133471e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.086612770817894, "grad_norm": 0.023925321438321916, "grad_rms": 2.0546121791758312e-05, "grad_to_weight": 0.0005968406853207359, "update_to_weight": 8.2794687554163e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200498709090347, "grad_norm": 0.05145843762768152, "grad_rms": 8.021641215122017e-05, "grad_to_weight": 0.0018247350218347043, "update_to_weight": 6.581558328401434e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411210309004566, "grad_norm": 0.005780199183351083, "grad_rms": 0.0004920432303913995, "grad_to_weight": 0.0006872018379047221, "update_to_weight": 4.030594775850855e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3153106476998875, "grad_norm": 0.05902882418550377, "grad_rms": 0.001366495836326751, "grad_to_weight": 0.04487823791948219, "update_to_weight": 0.0001222959979355154, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45150}
{"step": 45200, "lr": 1.2249605380542258e-05, "window_steps": 50, "elapsed_seconds": 2434.261326086, "steps_per_second": 1.7253693985078817, "eta_seconds": 2782.0129440982855, "images_seen": 46284800, "images_per_second": 1766.778264072071, "effective_batch_size": 1024, "loss": 0.3546660542488098, "edge_loss": 0.18228524923324585, "diag_loss": 0.17238086462020874, "jvp_rms": 1.990078091621399, "grad_norm": 0.20335575938224792, "grad_norm_postclip": 0.20335575938224792, "clip_fraction": 0.0, "ema_updates": 45200, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.2906031534075737, "updates_nfe_5": 12, "loss_nfe_3": 0.3843949536482493, "updates_nfe_3": 9, "loss_nfe_1": 0.3987779915332794, "updates_nfe_1": 7, "loss_nfe_4": 0.36873822945814866, "updates_nfe_4": 13, "loss_nfe_2": 0.355718904071384, "updates_nfe_2": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49696368399396, "grad_norm": 0.009150380177960086, "grad_rms": 2.6645013307767682e-05, "grad_to_weight": 0.0005546705656410073, "update_to_weight": 5.549894427271385e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25491587302323, "grad_norm": 0.0049114071743913265, "grad_rms": 1.4301537966337971e-05, "grad_to_weight": 0.0003021490368056799, "update_to_weight": 5.1427714263845587e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062038636083119, "grad_norm": 0.0010238094659569409, "grad_rms": 1.8329129668853346e-05, "grad_to_weight": 0.0002022524005761322, "update_to_weight": 2.6078247590545803e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854828374926992, "grad_norm": 0.005380114585053679, "grad_rms": 1.7216388709163424e-05, "grad_to_weight": 0.0004185286981775231, "update_to_weight": 6.366842199419186e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.049147150522733, "grad_norm": 0.013581197153262965, "grad_rms": 0.00030898293395604544, "grad_to_weight": 0.0026898002273229283, "update_to_weight": 2.5825346359693074e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.639982198558968, "grad_norm": 0.043940338542041255, "grad_rms": 0.0001058806412271423, "grad_to_weight": 0.0026406481700350926, "update_to_weight": 6.340987836237684e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566741344212453, "grad_norm": 0.019931579661708308, "grad_rms": 2.749860977596008e-05, "grad_to_weight": 0.0008113237072202856, "update_to_weight": 6.173776488595989e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828631838205617, "grad_norm": 0.011451277702024067, "grad_rms": 1.390158435713142e-05, "grad_to_weight": 0.00037144942928776396, "update_to_weight": 5.28720745562304e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61918773136781, "grad_norm": 0.008060610129245919, "grad_rms": 7.236027749208085e-06, "grad_to_weight": 0.00023283648916876083, "update_to_weight": 5.88018982839128e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06356876884379, "grad_norm": 0.006686480977408444, "grad_rms": 5.378349975684962e-06, "grad_to_weight": 0.00018040575151061016, "update_to_weight": 6.164041501189432e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.667315225408856, "grad_norm": 0.002346653392588768, "grad_rms": 4.2396784278900944e-06, "grad_to_weight": 9.513209569607195e-05, "update_to_weight": 3.779710979441001e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32407954742488, "grad_norm": 0.006666899439355994, "grad_rms": 5.362599334194453e-06, "grad_to_weight": 0.00017862193844284544, "update_to_weight": 6.38225235533274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45575283896154, "grad_norm": 0.016366110280590814, "grad_rms": 7.265985910107083e-06, "grad_to_weight": 0.00025791373592438255, "update_to_weight": 7.300863514477556e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7689191759769, "grad_norm": 0.022842476028613973, "grad_rms": 1.7044393354854346e-05, "grad_to_weight": 0.0004499303195610021, "update_to_weight": 6.043907425571751e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.087007117624005, "grad_norm": 0.022144072840978052, "grad_rms": 1.9016455796809298e-05, "grad_to_weight": 0.0005524002521815241, "update_to_weight": 7.796321106970188e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20070215527942, "grad_norm": 0.033459734158918814, "grad_rms": 5.215898401700117e-05, "grad_to_weight": 0.001186485853248688, "update_to_weight": 6.26433399672681e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.410743164901511, "grad_norm": 0.003577445539215751, "grad_rms": 0.0003045323878691161, "grad_to_weight": 0.00042534238284015443, "update_to_weight": 4.5359135448045645e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.314938233233964, "grad_norm": 0.04353191588721572, "grad_rms": 0.001007748038827024, "grad_to_weight": 0.03310567354951203, "update_to_weight": 8.058094601833767e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45200}
{"step": 45250, "lr": 1.2203376570814944e-05, "window_steps": 50, "elapsed_seconds": 2461.717599749, "steps_per_second": 1.7264368587336483, "eta_seconds": 2751.3314350135884, "images_seen": 46336000, "images_per_second": 1767.8713433432558, "effective_batch_size": 1024, "loss": 0.36668503284454346, "edge_loss": 0.1875472217798233, "diag_loss": 0.17913779616355896, "jvp_rms": 2.0763280391693115, "grad_norm": 0.2088628113269806, "grad_norm_postclip": 0.20744559168815613, "clip_fraction": 0.019999999552965164, "ema_updates": 45250, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3798664450645447, "updates_nfe_4": 15, "loss_nfe_2": 0.36059391498565674, "updates_nfe_2": 10, "loss_nfe_5": 0.3721995875239372, "updates_nfe_5": 10, "loss_nfe_1": 0.3926568031311035, "updates_nfe_1": 6, "loss_nfe_3": 0.3280421429210239, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.497648129319163, "grad_norm": 0.4727549489275161, "grad_rms": 0.0013766162345721113, "grad_to_weight": 0.02865589963016299, "update_to_weight": 0.0001259620025804215, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.252755738390857, "grad_norm": 0.44628243472730916, "grad_rms": 0.0012995308589443817, "grad_to_weight": 0.02745887786113338, "update_to_weight": 0.0001106801446720149, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061805583281469, "grad_norm": 0.013923746396048431, "grad_rms": 0.0002492750474140808, "grad_to_weight": 0.002750746974960255, "update_to_weight": 5.521405603549197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854021777424068, "grad_norm": 0.17651725965892584, "grad_rms": 0.0005648559539246464, "grad_to_weight": 0.013732453757698528, "update_to_weight": 0.0001279894899193989, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.049441603202317, "grad_norm": 0.04574850887740685, "grad_rms": 0.0010408146157910106, "grad_to_weight": 0.009060112478257694, "update_to_weight": 3.337181690396559e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640103241526834, "grad_norm": 0.2660858280124874, "grad_rms": 0.0006411725313509265, "grad_to_weight": 0.015990635643921183, "update_to_weight": 9.775909427503835e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566530163962412, "grad_norm": 0.030150899034023987, "grad_rms": 4.159769677081057e-05, "grad_to_weight": 0.0012273161424421874, "update_to_weight": 8.374923791421123e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82831407538318, "grad_norm": 0.012061549156642743, "grad_rms": 1.4642439685932965e-05, "grad_to_weight": 0.00039124906821531483, "update_to_weight": 6.36276907399028e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.619646119810895, "grad_norm": 0.003807298279672085, "grad_rms": 3.417820184760223e-06, "grad_to_weight": 0.00010997507792239911, "update_to_weight": 6.80382939283669e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.063830604276454, "grad_norm": 0.0028246664477683025, "grad_rms": 2.2720538310064753e-06, "grad_to_weight": 7.621086114726603e-05, "update_to_weight": 7.28579090370871e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.667445127558178, "grad_norm": 0.0012617428372664787, "grad_rms": 2.27957989262412e-06, "grad_to_weight": 5.115012238770016e-05, "update_to_weight": 4.8402123684825804e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32425452150301, "grad_norm": 0.002945850565784231, "grad_rms": 2.36952970813624e-06, "grad_to_weight": 7.89259049792163e-05, "update_to_weight": 7.62310080473965e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45648384363, "grad_norm": 0.008970096975168543, "grad_rms": 3.982412260240247e-06, "grad_to_weight": 0.0001413582416144067, "update_to_weight": 8.435085308135057e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76912355211581, "grad_norm": 0.0378973000093665, "grad_rms": 2.827786653414568e-05, "grad_to_weight": 0.0007464635463021919, "update_to_weight": 7.573825019100653e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08718373501296, "grad_norm": 0.17363933962899986, "grad_rms": 0.00014911461185819467, "grad_to_weight": 0.004331542489410144, "update_to_weight": 0.00010457949443655323, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200678347822482, "grad_norm": 0.49997346032755957, "grad_rms": 0.0007793877740417949, "grad_to_weight": 0.017729128858567512, "update_to_weight": 0.0001067516151079217, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411205320224019, "grad_norm": 0.004755538838691044, "grad_rms": 0.00040481834937126265, "grad_to_weight": 0.0005653813761098853, "update_to_weight": 7.275141657021401e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3154102032300954, "grad_norm": 0.08116988533012037, "grad_rms": 0.001879053358578816, "grad_to_weight": 0.06170689958980188, "update_to_weight": 0.000145376139903138, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45250}
{"step": 45300, "lr": 1.215761577810484e-05, "window_steps": 50, "elapsed_seconds": 2489.1908540140003, "steps_per_second": 1.7274689857814394, "eta_seconds": 2720.743491596698, "images_seen": 46387200, "images_per_second": 1768.928241440194, "effective_batch_size": 1024, "loss": 0.3753933012485504, "edge_loss": 0.1958368569612503, "diag_loss": 0.17955633997917175, "jvp_rms": 2.367304563522339, "grad_norm": 0.2607904076576233, "grad_norm_postclip": 0.2515968680381775, "clip_fraction": 0.019999999552965164, "ema_updates": 45300, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.4049637869000435, "updates_nfe_5": 10, "loss_nfe_4": 0.331068400825773, "updates_nfe_4": 7, "loss_nfe_1": 0.4012598435084025, "updates_nfe_1": 15, "loss_nfe_2": 0.364276647567749, "updates_nfe_2": 6, "loss_nfe_3": 0.34983229512969655, "updates_nfe_3": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.495630783966327, "grad_norm": 0.02012635578794254, "grad_rms": 5.86059821971405e-05, "grad_to_weight": 0.0012201022229174322, "update_to_weight": 5.0852521667041765e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25531012998137, "grad_norm": 0.008755146590161602, "grad_rms": 2.5494131704844483e-05, "grad_to_weight": 0.0005386022487515367, "update_to_weight": 9.152266696507523e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06183346728605, "grad_norm": 0.0021754308290552293, "grad_rms": 3.8946459353259616e-05, "grad_to_weight": 0.00042977131569316667, "update_to_weight": 3.3805360056595256e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854820659370176, "grad_norm": 0.008352451057008364, "grad_rms": 2.6727877594131978e-05, "grad_to_weight": 0.0006497524374966732, "update_to_weight": 7.825842179528291e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0488319028032285, "grad_norm": 0.023119107398566548, "grad_rms": 0.0005259779056176768, "grad_to_weight": 0.004579100244104044, "update_to_weight": 2.362821552348367e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640184853333963, "grad_norm": 0.062041885442534306, "grad_rms": 0.0001494989531614854, "grad_to_weight": 0.0037284372733457846, "update_to_weight": 8.16396844076442e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.567209660996173, "grad_norm": 0.02627960312458265, "grad_rms": 3.625666222423546e-05, "grad_to_weight": 0.0010697023995486607, "update_to_weight": 0.00010260382825477802, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828502159692356, "grad_norm": 0.0162829297113588, "grad_rms": 1.9767097336543152e-05, "grad_to_weight": 0.0005281777761051396, "update_to_weight": 9.096611840717538e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.61946629122209, "grad_norm": 0.01230175046174091, "grad_rms": 1.104330891553964e-05, "grad_to_weight": 0.00035534200204756104, "update_to_weight": 0.00011254518864380278, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06418465553623, "grad_norm": 0.009836670119090615, "grad_rms": 7.912241831627406e-06, "grad_to_weight": 0.00026539556206374887, "update_to_weight": 0.00011503158849047933, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66774822991298, "grad_norm": 0.0036526244237031604, "grad_rms": 6.599165016557928e-06, "grad_to_weight": 0.0001480728759536251, "update_to_weight": 7.794494713943273e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32517353076289, "grad_norm": 0.009758999558879098, "grad_rms": 7.849766598835081e-06, "grad_to_weight": 0.00026145891996552585, "update_to_weight": 0.00010830585609897219, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45687819791484, "grad_norm": 0.020675155239368553, "grad_rms": 9.179052571623354e-06, "grad_to_weight": 0.00032581425097662507, "update_to_weight": 0.00011363931776811221, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76937601562211, "grad_norm": 0.027200684341117574, "grad_rms": 2.0296362042822134e-05, "grad_to_weight": 0.0005357695224134275, "update_to_weight": 8.101459315959834e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08748367827856, "grad_norm": 0.029086689180520132, "grad_rms": 2.4978500705313884e-05, "grad_to_weight": 0.0007255803186339876, "update_to_weight": 8.689826729924543e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200718387624896, "grad_norm": 0.05212259625663091, "grad_rms": 8.125174133667011e-05, "grad_to_weight": 0.001848271931948495, "update_to_weight": 6.37902223922292e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411084228006816, "grad_norm": 0.006180736095531882, "grad_rms": 0.0005261391966217843, "grad_to_weight": 0.000734832267515711, "update_to_weight": 4.419084048024469e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.315257309661838, "grad_norm": 0.06928229607604078, "grad_rms": 0.0016038599857849927, "grad_to_weight": 0.052675849483667755, "update_to_weight": 0.0001460905573621956, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45300}
{"step": 45350, "lr": 1.2112323502844616e-05, "window_steps": 50, "elapsed_seconds": 2516.693489181, "steps_per_second": 1.7284583993641622, "eta_seconds": 2690.2585574003792, "images_seen": 46438400, "images_per_second": 1769.9414009489021, "effective_batch_size": 1024, "loss": 0.3599410057067871, "edge_loss": 0.18573828041553497, "diag_loss": 0.1742027848958969, "jvp_rms": 2.015429735183716, "grad_norm": 0.19853398203849792, "grad_norm_postclip": 0.19853398203849792, "clip_fraction": 0.0, "ema_updates": 45350, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.35291182808578014, "updates_nfe_3": 16, "loss_nfe_5": 0.3469076305627823, "updates_nfe_5": 8, "loss_nfe_2": 0.3911743760108948, "updates_nfe_2": 10, "loss_nfe_1": 0.3980423543188307, "updates_nfe_1": 9, "loss_nfe_4": 0.29729658365249634, "updates_nfe_4": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49780536274778, "grad_norm": 0.04891110759824712, "grad_rms": 0.0001424243678958796, "grad_to_weight": 0.0029647038816865254, "update_to_weight": 4.5531937149103865e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25227692131023, "grad_norm": 0.015234194962342464, "grad_rms": 4.4360487718580666e-05, "grad_to_weight": 0.0009373575798703724, "update_to_weight": 4.186143635056705e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061790133969558, "grad_norm": 0.0014302442618980328, "grad_rms": 2.5605479736368286e-05, "grad_to_weight": 0.0002825570053368464, "update_to_weight": 2.5406810922431547e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854354749856418, "grad_norm": 0.014676850300849704, "grad_rms": 4.696598107921331e-05, "grad_to_weight": 0.001141780399441181, "update_to_weight": 6.0538010193363564e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.049079153877866, "grad_norm": 0.013649528411644985, "grad_rms": 0.0003105375239128456, "grad_to_weight": 0.0027033698612472096, "update_to_weight": 2.0180913436969496e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640664429237326, "grad_norm": 0.056720599079992266, "grad_rms": 0.0001366765391584588, "grad_to_weight": 0.0034085537462275406, "update_to_weight": 6.457090306911547e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.567317732832322, "grad_norm": 0.015453731764456377, "grad_rms": 2.1320745599986356e-05, "grad_to_weight": 0.0006290361826437267, "update_to_weight": 7.427136965282184e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828604120706633, "grad_norm": 0.007624097415705147, "grad_rms": 9.255476648922684e-06, "grad_to_weight": 0.0002473059560482751, "update_to_weight": 6.403651588476204e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62056464146346, "grad_norm": 0.00469531184279924, "grad_rms": 4.214991947372487e-06, "grad_to_weight": 0.00013562204693726688, "update_to_weight": 7.58463253460062e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06445307353968, "grad_norm": 0.003970736826885811, "grad_rms": 3.1939090813969347e-06, "grad_to_weight": 0.00010713059272741623, "update_to_weight": 8.090481375205803e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.667464922111346, "grad_norm": 0.0013350264889192237, "grad_rms": 2.411980833474784e-06, "grad_to_weight": 5.412094405057963e-05, "update_to_weight": 5.534263837278768e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32541227323802, "grad_norm": 0.00408192749952495, "grad_rms": 3.283346612664225e-06, "grad_to_weight": 0.00010936054690149142, "update_to_weight": 8.331184964683092e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45734757250337, "grad_norm": 0.010582982285114931, "grad_rms": 4.698477454459755e-06, "grad_to_weight": 0.0001667731585065593, "update_to_weight": 8.97257890169956e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76983044676976, "grad_norm": 0.018536520364839014, "grad_rms": 1.3831414078439336e-05, "grad_to_weight": 0.00036510896730832795, "update_to_weight": 6.89984904350709e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08780797989349, "grad_norm": 0.025196884858430378, "grad_rms": 2.1638090272217133e-05, "grad_to_weight": 0.0006285423456196002, "update_to_weight": 8.357828290526924e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200984596757653, "grad_norm": 0.052073724674968284, "grad_rms": 8.11755574663883e-05, "grad_to_weight": 0.0018465215104920608, "update_to_weight": 6.656965495059008e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411474256628376, "grad_norm": 0.005417789126061842, "grad_rms": 0.0004611928375833876, "grad_to_weight": 0.0006440950730833584, "update_to_weight": 4.157522322202371e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3155886321126722, "grad_norm": 0.06413417708248297, "grad_rms": 0.0014846829012558546, "grad_to_weight": 0.04874941567372122, "update_to_weight": 0.00010589747104379204, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45350}
{"step": 45400, "lr": 1.2067500240343352e-05, "window_steps": 50, "elapsed_seconds": 2544.1975384380003, "steps_per_second": 1.7294254606901955, "eta_seconds": 2659.842881094273, "images_seen": 46489600, "images_per_second": 1770.9316717467602, "effective_batch_size": 1024, "loss": 0.3821423351764679, "edge_loss": 0.19546760618686676, "diag_loss": 0.18667466938495636, "jvp_rms": 2.1219139099121094, "grad_norm": 0.31918227672576904, "grad_norm_postclip": 0.28525328636169434, "clip_fraction": 0.05999999865889549, "ema_updates": 45400, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3618821104367574, "updates_nfe_2": 9, "loss_nfe_5": 0.42931244235772353, "updates_nfe_5": 13, "loss_nfe_3": 0.3495686839024226, "updates_nfe_3": 15, "loss_nfe_1": 0.4023421009381612, "updates_nfe_1": 6, "loss_nfe_4": 0.3730757364204952, "updates_nfe_4": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.497509392696504, "grad_norm": 0.0061429004724855634, "grad_rms": 1.7887526163329232e-05, "grad_to_weight": 0.00037235320352083226, "update_to_weight": 2.9291293735313092e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.254891466325798, "grad_norm": 0.012638694801775482, "grad_rms": 3.680264476849269e-05, "grad_to_weight": 0.0007775317865368861, "update_to_weight": 4.3679566060847e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061747177098367, "grad_norm": 0.003191470494278104, "grad_rms": 5.713648727526328e-05, "grad_to_weight": 0.0006305076849191046, "update_to_weight": 2.2919858119880525e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854878822684345, "grad_norm": 0.010207475547175703, "grad_rms": 3.266396356086236e-05, "grad_to_weight": 0.000794054591099147, "update_to_weight": 4.874260200726052e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.048820002713486, "grad_norm": 0.018971885922782167, "grad_rms": 0.00043162535002978385, "grad_to_weight": 0.003757687125424502, "update_to_weight": 1.7078268852916517e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640480119674468, "grad_norm": 0.03757917736536822, "grad_rms": 9.055249750583416e-05, "grad_to_weight": 0.0022582988648829545, "update_to_weight": 5.222554773285483e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56733388145667, "grad_norm": 0.02329971747716015, "grad_rms": 3.214546210933036e-05, "grad_to_weight": 0.0009484023618349034, "update_to_weight": 7.223500721192374e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828734788697947, "grad_norm": 0.013992865955934585, "grad_rms": 1.6987013287615422e-05, "grad_to_weight": 0.0004538903737646891, "update_to_weight": 6.91290222219985e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.620693338328195, "grad_norm": 0.010395666681601416, "grad_rms": 9.332213241111441e-06, "grad_to_weight": 0.0003002732088583704, "update_to_weight": 8.667935616331134e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.064780771411726, "grad_norm": 0.008389145071879687, "grad_rms": 6.747907957236022e-06, "grad_to_weight": 0.00022633737195473395, "update_to_weight": 9.216192897034869e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66778286994496, "grad_norm": 0.0033012133148370903, "grad_rms": 5.9642735995784365e-06, "grad_to_weight": 0.00013382691635652685, "update_to_weight": 6.286602355839496e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32586522565485, "grad_norm": 0.008069626681108106, "grad_rms": 6.490899564474051e-06, "grad_to_weight": 0.00021619396180966977, "update_to_weight": 9.141142178541567e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.457772699769606, "grad_norm": 0.019816744650959763, "grad_rms": 8.797948012657873e-06, "grad_to_weight": 0.00031228238571682034, "update_to_weight": 0.00010053732563221327, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.769888152021274, "grad_norm": 0.025411456301709524, "grad_rms": 1.896129195379133e-05, "grad_to_weight": 0.0005005222037444617, "update_to_weight": 7.214086599810538e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08808203555095, "grad_norm": 0.025853879316160483, "grad_rms": 2.2202291182948123e-05, "grad_to_weight": 0.0006449268212241415, "update_to_weight": 7.404892369974661e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200857985410323, "grad_norm": 0.04180966119295481, "grad_rms": 6.517533700542797e-05, "grad_to_weight": 0.0014825669919186495, "update_to_weight": 4.809070680682019e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411576296015593, "grad_norm": 0.005211957265100958, "grad_rms": 0.00044367126599526973, "grad_to_weight": 0.0006196171896544249, "update_to_weight": 3.313035126246182e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3158081235316816, "grad_norm": 0.06536367986948846, "grad_rms": 0.0015131454441300823, "grad_to_weight": 0.049675692603302776, "update_to_weight": 4.6073036096148034e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45400}
{"step": 45450, "lr": 1.2023146480781069e-05, "window_steps": 50, "elapsed_seconds": 2571.700454292, "steps_per_second": 1.7303725994111177, "eta_seconds": 2629.4914757367637, "images_seen": 46540800, "images_per_second": 1771.9015417969845, "effective_batch_size": 1024, "loss": 0.37109389901161194, "edge_loss": 0.1917327493429184, "diag_loss": 0.17936113476753235, "jvp_rms": 2.2851171493530273, "grad_norm": 0.15721982717514038, "grad_norm_postclip": 0.15721982717514038, "clip_fraction": 0.0, "ema_updates": 45450, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3512862351807681, "updates_nfe_4": 11, "loss_nfe_3": 0.3502182479088123, "updates_nfe_3": 13, "loss_nfe_2": 0.37314479218588936, "updates_nfe_2": 9, "loss_nfe_5": 0.401617237499782, "updates_nfe_5": 7, "loss_nfe_1": 0.39680867791175845, "updates_nfe_1": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49864422401031, "grad_norm": 0.011275029407261192, "grad_rms": 3.283178433673785e-05, "grad_to_weight": 0.0006833912686505935, "update_to_weight": 5.00828685028348e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.254738454272246, "grad_norm": 0.004645904671290202, "grad_rms": 1.3528420610469575e-05, "grad_to_weight": 0.00028581848206047977, "update_to_weight": 3.821076705376199e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061182863294707, "grad_norm": 0.0006945712657606811, "grad_rms": 1.243482036229049e-05, "grad_to_weight": 0.0001372349675009632, "update_to_weight": 2.253215344794554e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854661005655206, "grad_norm": 0.005647543352322405, "grad_rms": 1.807216185981368e-05, "grad_to_weight": 0.00043933817856712495, "update_to_weight": 5.2364470483733016e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0487922357283415, "grad_norm": 0.009737857815733587, "grad_rms": 0.00022154393639954502, "grad_to_weight": 0.0019287499586183306, "update_to_weight": 2.402588427748364e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640539722492758, "grad_norm": 0.05000134052774376, "grad_rms": 0.00012048550768967892, "grad_to_weight": 0.0030047907917408307, "update_to_weight": 6.757823626986618e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.566998484771826, "grad_norm": 0.015433594768785305, "grad_rms": 2.1292963587952314e-05, "grad_to_weight": 0.0006282246802901876, "update_to_weight": 7.963232829731056e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828690242854037, "grad_norm": 0.006900429386873048, "grad_rms": 8.376960520753837e-06, "grad_to_weight": 0.00022383141588288978, "update_to_weight": 6.783089316212279e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62093310107431, "grad_norm": 0.0037611476278740018, "grad_rms": 3.376390641375743e-06, "grad_to_weight": 0.00010863796238228169, "update_to_weight": 8.029103130520757e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06484499326822, "grad_norm": 0.0032594919580529896, "grad_rms": 2.6218108677150823e-06, "grad_to_weight": 8.794025602009086e-05, "update_to_weight": 7.973863373894987e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66820720638903, "grad_norm": 0.001102823996276971, "grad_rms": 1.99246259440854e-06, "grad_to_weight": 4.470628882958877e-05, "update_to_weight": 5.0535472431150085e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.326509487854274, "grad_norm": 0.003353200131409838, "grad_rms": 2.6971861441270146e-06, "grad_to_weight": 8.983428071411125e-05, "update_to_weight": 7.82708798402289e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.458726821183745, "grad_norm": 0.009144631797390721, "grad_rms": 4.059899674008523e-06, "grad_to_weight": 0.0001441036127806155, "update_to_weight": 8.102937931454356e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76988334325282, "grad_norm": 0.01857014141121305, "grad_rms": 1.3856501128489561e-05, "grad_to_weight": 0.0003657708111255878, "update_to_weight": 6.256641495940268e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08847484539275, "grad_norm": 0.021099201961451614, "grad_rms": 1.8119161923339084e-05, "grad_to_weight": 0.0005263159060758452, "update_to_weight": 7.421354286501831e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200748687978127, "grad_norm": 0.033397052133839485, "grad_rms": 5.206127162249326e-05, "grad_to_weight": 0.0011842611876498343, "update_to_weight": 5.348621846091029e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411722777060891, "grad_norm": 0.0027338793018685955, "grad_rms": 0.00023272326100218102, "grad_to_weight": 0.00032500825031038765, "update_to_weight": 4.0520631875187375e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3160520807268554, "grad_norm": 0.03746668806280484, "grad_rms": 0.0008673402180243522, "grad_to_weight": 0.028469001046001155, "update_to_weight": 6.980589289632392e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45450}
{"step": 45500, "lr": 1.1979262709203377e-05, "window_steps": 50, "elapsed_seconds": 2599.185082522, "steps_per_second": 1.7313118755027754, "eta_seconds": 2599.185082522, "images_seen": 46592000, "images_per_second": 1772.863360514842, "effective_batch_size": 1024, "loss": 0.35728999972343445, "edge_loss": 0.18617993593215942, "diag_loss": 0.17111004889011383, "jvp_rms": 2.0214061737060547, "grad_norm": 0.20125561952590942, "grad_norm_postclip": 0.20125561952590942, "clip_fraction": 0.0, "ema_updates": 45500, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3550542157429915, "updates_nfe_2": 13, "loss_nfe_1": 0.4040657336061651, "updates_nfe_1": 11, "loss_nfe_5": 0.321593114733696, "updates_nfe_5": 10, "loss_nfe_3": 0.3567018136382103, "updates_nfe_3": 10, "loss_nfe_4": 0.33685359358787537, "updates_nfe_4": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.500064734129644, "grad_norm": 0.17400973723114402, "grad_rms": 0.0005066993582815936, "grad_to_weight": 0.010546003305745382, "update_to_weight": 4.7130777784162716e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2541216978331, "grad_norm": 0.14282032239847237, "grad_rms": 0.0004158788287390398, "grad_to_weight": 0.008786714228767728, "update_to_weight": 5.1442397330255835e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0619088286920215, "grad_norm": 0.005477190303100916, "grad_rms": 9.805743609987789e-05, "grad_to_weight": 0.0010820404887687796, "update_to_weight": 2.1840102105877523e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.855962511407949, "grad_norm": 0.05891128963697786, "grad_rms": 0.00018851636813943477, "grad_to_weight": 0.004582409880606136, "update_to_weight": 6.034125019321976e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.048443340934508, "grad_norm": 0.0304287715043746, "grad_rms": 0.0006922785222833537, "grad_to_weight": 0.006027357236565913, "update_to_weight": 2.443536808641433e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640772629920257, "grad_norm": 0.17167636924548693, "grad_rms": 0.00041367919916839545, "grad_to_weight": 0.010316610476175325, "update_to_weight": 6.769981387132315e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.567510273431377, "grad_norm": 0.021725026798902227, "grad_rms": 2.99729396492844e-05, "grad_to_weight": 0.0008842990826952797, "update_to_weight": 7.237401761643076e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8287387483254, "grad_norm": 0.008288303983212926, "grad_rms": 1.0061807948279614e-05, "grad_to_weight": 0.0002688499212009814, "update_to_weight": 5.643172122354533e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62072507172817, "grad_norm": 0.002878229850040518, "grad_rms": 2.583793376623787e-06, "grad_to_weight": 8.313603611932801e-05, "update_to_weight": 6.271643212059484e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.064604572413494, "grad_norm": 0.0026250876027753382, "grad_rms": 2.1115202289904303e-06, "grad_to_weight": 7.082464882760797e-05, "update_to_weight": 6.664244947133746e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668490505665826, "grad_norm": 0.0009448819765999485, "grad_rms": 1.7071101108262275e-06, "grad_to_weight": 3.8303193962473276e-05, "update_to_weight": 4.188981895545681e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32709323867099, "grad_norm": 0.0025708657585483867, "grad_rms": 2.0679062479494427e-06, "grad_to_weight": 6.887398764511782e-05, "update_to_weight": 7.046620947841367e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45932313978282, "grad_norm": 0.007350018875917802, "grad_rms": 3.263153716786022e-06, "grad_to_weight": 0.0001158225224011262, "update_to_weight": 7.975458198743009e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.76989296078927, "grad_norm": 0.027992985680086804, "grad_rms": 2.0887554331261135e-05, "grad_to_weight": 0.0005513697990599747, "update_to_weight": 7.068915767234741e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.089070142893995, "grad_norm": 0.06276923913333993, "grad_rms": 5.3903745257269954e-05, "grad_to_weight": 0.0015657444512832163, "update_to_weight": 8.952102935928689e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200861231862227, "grad_norm": 0.18709089430806666, "grad_rms": 0.00029164819181146264, "grad_to_weight": 0.0066342262659938, "update_to_weight": 7.717845711558963e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411712346606414, "grad_norm": 0.006192597944288326, "grad_rms": 0.0005271489442438638, "grad_to_weight": 0.0007361875548189236, "update_to_weight": 6.1804230847120934e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3160936114517097, "grad_norm": 0.115131546467314, "grad_rms": 0.002665253476556298, "grad_to_weight": 0.08747975483318302, "update_to_weight": 0.00022910723354147743, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45500}
{"step": 45550, "lr": 1.1935849405516155e-05, "window_steps": 50, "elapsed_seconds": 2628.9243598890002, "steps_per_second": 1.7307458782085736, "eta_seconds": 2571.1458025288025, "images_seen": 46643200, "images_per_second": 1772.2837792855794, "effective_batch_size": 1024, "loss": 0.35112711787223816, "edge_loss": 0.18134161829948425, "diag_loss": 0.1697855144739151, "jvp_rms": 1.7677342891693115, "grad_norm": 0.2827598452568054, "grad_norm_postclip": 0.2827598452568054, "clip_fraction": 0.0, "ema_updates": 45550, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3155979730866172, "updates_nfe_4": 11, "loss_nfe_2": 0.3509090624072335, "updates_nfe_2": 11, "loss_nfe_5": 0.2905105969735554, "updates_nfe_5": 14, "loss_nfe_3": 0.45857784897089005, "updates_nfe_3": 8, "loss_nfe_1": 0.4148353387912114, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.498216937274616, "grad_norm": 0.0907231446859723, "grad_rms": 0.0002641769358723151, "grad_to_weight": 0.005498966647783642, "update_to_weight": 5.911762195171112e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25341948697354, "grad_norm": 0.03833672745244283, "grad_rms": 0.00011163280577204641, "grad_to_weight": 0.0023586868894369065, "update_to_weight": 7.051166172461818e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.061932378901297, "grad_norm": 0.0030653608654462516, "grad_rms": 5.487876275111124e-05, "grad_to_weight": 0.0006055712791073663, "update_to_weight": 3.203354308235994e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.854577913717888, "grad_norm": 0.028532454508323857, "grad_rms": 9.13039712957944e-05, "grad_to_weight": 0.002219633713361772, "update_to_weight": 6.962591793843697e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.048693822258692, "grad_norm": 0.02878327792941552, "grad_rms": 0.0006548422471995681, "grad_to_weight": 0.0057011335887543315, "update_to_weight": 2.53576190055397e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640466365147624, "grad_norm": 0.10662567536921672, "grad_rms": 0.0002569300841541791, "grad_to_weight": 0.006407613406348831, "update_to_weight": 7.520638765749463e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56714009726394, "grad_norm": 0.017279383636992822, "grad_rms": 2.3839506745952027e-05, "grad_to_weight": 0.0007033534863472871, "update_to_weight": 8.372878955390162e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828945638152405, "grad_norm": 0.007936335330668743, "grad_rms": 9.634526203680582e-06, "grad_to_weight": 0.00025743129277983225, "update_to_weight": 6.533586367673416e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.6235404137308, "grad_norm": 0.00309834431520956, "grad_rms": 2.781390624527686e-06, "grad_to_weight": 8.948664053953411e-05, "update_to_weight": 7.024190121140873e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06579184105862, "grad_norm": 0.0024397159663272134, "grad_rms": 1.962414363027155e-06, "grad_to_weight": 6.582122882438153e-05, "update_to_weight": 7.450820996374193e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668428650206167, "grad_norm": 0.0009766790247851998, "grad_rms": 1.7645575633078543e-06, "grad_to_weight": 3.9592267453850865e-05, "update_to_weight": 4.8677496337835423e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32796476001605, "grad_norm": 0.0025747768044026755, "grad_rms": 2.0710521438917754e-06, "grad_to_weight": 6.897715482095222e-05, "update_to_weight": 8.043202129704775e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.46152562281339, "grad_norm": 0.007561446645743479, "grad_rms": 3.357020321020054e-06, "grad_to_weight": 0.00011915009246208952, "update_to_weight": 9.076020991387142e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.770530118520774, "grad_norm": 0.025424967291472995, "grad_rms": 1.897137346263711e-05, "grad_to_weight": 0.0005007819936510398, "update_to_weight": 7.774492325908109e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.08943706099136, "grad_norm": 0.036752379646305636, "grad_rms": 3.156149313590571e-05, "grad_to_weight": 0.0009167596838636377, "update_to_weight": 9.318763446087251e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.200789809834006, "grad_norm": 0.10821742762017364, "grad_rms": 0.00016869563430458554, "grad_to_weight": 0.0038373899578669503, "update_to_weight": 7.969598009786474e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412170820895275, "grad_norm": 0.011893587419493194, "grad_rms": 0.0010124493964993673, "grad_to_weight": 0.0014138547198721061, "update_to_weight": 6.106726644848834e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3165274997190448, "grad_norm": 0.16298979935073346, "grad_rms": 0.0037731546452047806, "grad_to_weight": 0.12380280653880492, "update_to_weight": 0.00014266291107118036, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45550}
{"step": 45600, "lr": 1.189290704448035e-05, "window_steps": 50, "elapsed_seconds": 2656.415047795, "steps_per_second": 1.7316571082588557, "eta_seconds": 2540.9187413691307, "images_seen": 46694400, "images_per_second": 1773.2168788570682, "effective_batch_size": 1024, "loss": 0.36202991008758545, "edge_loss": 0.1891690343618393, "diag_loss": 0.17286084592342377, "jvp_rms": 2.002695322036743, "grad_norm": 0.2661590278148651, "grad_norm_postclip": 0.24568799138069153, "clip_fraction": 0.019999999552965164, "ema_updates": 45600, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3590856481056947, "updates_nfe_3": 13, "loss_nfe_5": 0.31926092963952285, "updates_nfe_5": 13, "loss_nfe_4": 0.384843073785305, "updates_nfe_4": 4, "loss_nfe_2": 0.3636505603790283, "updates_nfe_2": 9, "loss_nfe_1": 0.4064329645850442, "updates_nfe_1": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49872468585877, "grad_norm": 0.023268138648593367, "grad_rms": 6.775454701128793e-05, "grad_to_weight": 0.0014102992256447998, "update_to_weight": 7.685396187926415e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25485485621094, "grad_norm": 0.006722836301393311, "grad_rms": 1.9576242651428532e-05, "grad_to_weight": 0.0004135894390237838, "update_to_weight": 0.00010714032553887558, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062171077639229, "grad_norm": 0.0007944894112824973, "grad_rms": 1.422364211715572e-05, "grad_to_weight": 0.00015694637717637386, "update_to_weight": 3.342565941478628e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.857234809183371, "grad_norm": 0.009570740184180486, "grad_rms": 3.0626407791204614e-05, "grad_to_weight": 0.0007443855794983629, "update_to_weight": 0.00010164169587444644, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.047935917600207, "grad_norm": 0.02913950536356491, "grad_rms": 0.0006629467019480707, "grad_to_weight": 0.005772558495040851, "update_to_weight": 2.880047696571385e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.640945015995655, "grad_norm": 0.08124250459943871, "grad_rms": 0.000195765639667464, "grad_to_weight": 0.0048820847927414325, "update_to_weight": 7.7675205736063e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.567793491699273, "grad_norm": 0.015701281896944308, "grad_rms": 2.16622782134977e-05, "grad_to_weight": 0.0006391002066282144, "update_to_weight": 9.293805884340692e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82880210229551, "grad_norm": 0.006828849653794191, "grad_rms": 8.290064392343573e-06, "grad_to_weight": 0.00022150875765898525, "update_to_weight": 7.013391461456869e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62196794034822, "grad_norm": 0.003759626557825036, "grad_rms": 3.375025173389285e-06, "grad_to_weight": 0.00010859078156107907, "update_to_weight": 7.991610836581288e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.065378524389175, "grad_norm": 0.003148825648466704, "grad_rms": 2.532795114064762e-06, "grad_to_weight": 8.495328454273746e-05, "update_to_weight": 8.231265633924664e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668390299743287, "grad_norm": 0.0010697023061996852, "grad_rms": 1.9326219228549896e-06, "grad_to_weight": 4.3363279614187765e-05, "update_to_weight": 5.356715309073882e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.328149526467456, "grad_norm": 0.003243892031128463, "grad_rms": 2.609262882178491e-06, "grad_to_weight": 8.690203163776943e-05, "update_to_weight": 8.481534292101444e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.45989829338092, "grad_norm": 0.00899244429093127, "grad_rms": 3.992333694147071e-06, "grad_to_weight": 0.00014170278447907963, "update_to_weight": 9.378435699794688e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7702055297322, "grad_norm": 0.022567807056282748, "grad_rms": 1.6839443331007433e-05, "grad_to_weight": 0.00044450887722064707, "update_to_weight": 8.160293272852897e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.089756779417016, "grad_norm": 0.03612402377337816, "grad_rms": 3.102188591152561e-05, "grad_to_weight": 0.000901078646401892, "update_to_weight": 0.00010325300777842033, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201396891308306, "grad_norm": 0.08348118219458535, "grad_rms": 0.0001301353330282538, "grad_to_weight": 0.0029601789768192058, "update_to_weight": 8.292079847859748e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411474256628376, "grad_norm": 0.011771318871942462, "grad_rms": 0.0010020412065384727, "grad_to_weight": 0.0013994358792296696, "update_to_weight": 9.425339482964865e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.315823978046953, "grad_norm": 0.11819789733092602, "grad_rms": 0.0027362383851269787, "grad_to_weight": 0.0898280463822862, "update_to_weight": 0.0001239481138255038, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45600}
{"step": 45650, "lr": 1.1850436095706742e-05, "window_steps": 50, "elapsed_seconds": 2683.89216141, "steps_per_second": 1.7325584339264557, "eta_seconds": 2510.7378284158067, "images_seen": 46745600, "images_per_second": 1774.1398363406906, "effective_batch_size": 1024, "loss": 0.3500475585460663, "edge_loss": 0.18234717845916748, "diag_loss": 0.1677003651857376, "jvp_rms": 1.7905077934265137, "grad_norm": 0.24329276382923126, "grad_norm_postclip": 0.24329276382923126, "clip_fraction": 0.0, "ema_updates": 45650, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39989741891622543, "updates_nfe_1": 12, "loss_nfe_5": 0.2905326377261769, "updates_nfe_5": 11, "loss_nfe_4": 0.39301967322826387, "updates_nfe_4": 10, "loss_nfe_3": 0.29862486322720844, "updates_nfe_3": 9, "loss_nfe_2": 0.36124101281166077, "updates_nfe_2": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.497958895120274, "grad_norm": 0.023096701336312163, "grad_rms": 6.725533830319696e-05, "grad_to_weight": 0.0013999732623375398, "update_to_weight": 5.564498012896551e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2539564744778, "grad_norm": 0.00959862468790259, "grad_rms": 2.79502575380915e-05, "grad_to_weight": 0.0005905408140457674, "update_to_weight": 4.745226055503367e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062744511480447, "grad_norm": 0.0008430790099888389, "grad_rms": 1.5093535476086179e-05, "grad_to_weight": 0.0001665260824592363, "update_to_weight": 2.63833985971942e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.855261032202135, "grad_norm": 0.009130842672934531, "grad_rms": 2.9218733953393894e-05, "grad_to_weight": 0.0007102806119659475, "update_to_weight": 5.69297041888306e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.048392147520484, "grad_norm": 0.01879189567806775, "grad_rms": 0.00042753043017347607, "grad_to_weight": 0.0037223526083046826, "update_to_weight": 2.6274541251216683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64050671172747, "grad_norm": 0.05167279368328789, "grad_rms": 0.00012451311734773418, "grad_to_weight": 0.003105241599816866, "update_to_weight": 6.69698908486052e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.567107799749877, "grad_norm": 0.01783911352853626, "grad_rms": 2.461173825638587e-05, "grad_to_weight": 0.0007261381223197092, "update_to_weight": 8.492301264402871e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828895153203067, "grad_norm": 0.010196381809485787, "grad_rms": 1.2378169978100586e-05, "grad_to_weight": 0.00033074107128443106, "update_to_weight": 7.517786485837055e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.624488799249804, "grad_norm": 0.006802730231919753, "grad_rms": 6.106826150783344e-06, "grad_to_weight": 0.00019647164385188398, "update_to_weight": 8.876172160071009e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06703999649233, "grad_norm": 0.005661616717086423, "grad_rms": 4.553988299011246e-06, "grad_to_weight": 0.00015273991982154992, "update_to_weight": 9.474690919959089e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66884184173428, "grad_norm": 0.0019025687884221462, "grad_rms": 3.4373546069161294e-06, "grad_to_weight": 7.71243660577294e-05, "update_to_weight": 6.226689016709276e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.329735537148366, "grad_norm": 0.005700990721701884, "grad_rms": 4.585659244831878e-06, "grad_to_weight": 0.00015271982615651248, "update_to_weight": 9.31789014378622e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.46281437656031, "grad_norm": 0.014440595827999834, "grad_rms": 6.411124208556301e-06, "grad_to_weight": 0.00022754420789339276, "update_to_weight": 9.82425262821222e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77098694367064, "grad_norm": 0.023043047558865867, "grad_rms": 1.719405401568271e-05, "grad_to_weight": 0.00045386250979189454, "update_to_weight": 7.671937175470678e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09032160907106, "grad_norm": 0.028435506726258948, "grad_rms": 2.441929091377995e-05, "grad_to_weight": 0.0007092860716743408, "update_to_weight": 8.767104328375852e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201057100435637, "grad_norm": 0.06293722049609402, "grad_rms": 9.811020799922343e-05, "grad_to_weight": 0.0022317326712948567, "update_to_weight": 7.043432438624016e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412091462643454, "grad_norm": 0.00797513715168679, "grad_rms": 0.0006788887584070105, "grad_to_weight": 0.0009480564003735459, "update_to_weight": 6.8918693576132185e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.316538908751825, "grad_norm": 0.08494015143426555, "grad_rms": 0.001966333649254562, "grad_to_weight": 0.06451776766308791, "update_to_weight": 0.00010597785553901515, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45650}
{"step": 45700, "lr": 1.1808437023650818e-05, "window_steps": 50, "elapsed_seconds": 2711.370447855, "steps_per_second": 1.7334407416434854, "eta_seconds": 2480.6155161226593, "images_seen": 46796800, "images_per_second": 1775.043319442929, "effective_batch_size": 1024, "loss": 0.3743952214717865, "edge_loss": 0.1935236006975174, "diag_loss": 0.18087156116962433, "jvp_rms": 2.1811327934265137, "grad_norm": 0.22240397334098816, "grad_norm_postclip": 0.22240397334098816, "clip_fraction": 0.0, "ema_updates": 45700, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.40218834578990936, "updates_nfe_5": 10, "loss_nfe_4": 0.3560766130685806, "updates_nfe_4": 10, "loss_nfe_2": 0.35493939369916916, "updates_nfe_2": 12, "loss_nfe_1": 0.3999451498190562, "updates_nfe_1": 9, "loss_nfe_3": 0.36425892843140495, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.500291301583417, "grad_norm": 0.11995197968248619, "grad_rms": 0.0003492884484331275, "grad_to_weight": 0.007269688606708128, "update_to_weight": 5.1656687314514115e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25446903691028, "grad_norm": 0.03838942422471656, "grad_rms": 0.00011178625362570959, "grad_to_weight": 0.0023617765758784694, "update_to_weight": 5.6805902706335784e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062350047210118, "grad_norm": 0.0036929653540732944, "grad_rms": 6.611468548410467e-05, "grad_to_weight": 0.0007294962457423313, "update_to_weight": 2.4660085199043442e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.85641352689054, "grad_norm": 0.02131878302302182, "grad_rms": 6.822019299557275e-05, "grad_to_weight": 0.0016582216322173634, "update_to_weight": 5.8757101387404426e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.047974080026766, "grad_norm": 0.06812970544374636, "grad_rms": 0.0015500044686791326, "grad_to_weight": 0.013496445180515886, "update_to_weight": 2.3117414082702818e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641372305696333, "grad_norm": 0.14446774582552027, "grad_rms": 0.00034811606082667436, "grad_to_weight": 0.008681239934525655, "update_to_weight": 5.84352024333407e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56782081960504, "grad_norm": 0.03632752209081054, "grad_rms": 5.011927659812736e-05, "grad_to_weight": 0.0014786627742669508, "update_to_weight": 7.208289834486054e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.828832789327976, "grad_norm": 0.022749973077925142, "grad_rms": 2.7617937324962835e-05, "grad_to_weight": 0.0007379446777433788, "update_to_weight": 6.337803888637215e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62387358558127, "grad_norm": 0.015869346947429128, "grad_rms": 1.4245948263490865e-05, "grad_to_weight": 0.0004583353999431693, "update_to_weight": 7.523517917754608e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06695931225341, "grad_norm": 0.012643759850077674, "grad_rms": 1.0170157622820052e-05, "grad_to_weight": 0.0003411059359783516, "update_to_weight": 8.084347104185628e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668542464132077, "grad_norm": 0.0047997353009465, "grad_rms": 8.671640336520502e-06, "grad_to_weight": 0.00019456906738309684, "update_to_weight": 5.225734276911526e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.32952788803724, "grad_norm": 0.011915810058188666, "grad_rms": 9.584622606907457e-06, "grad_to_weight": 0.0003192060208724807, "update_to_weight": 8.26654750048271e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.463256778528944, "grad_norm": 0.024977698567702156, "grad_rms": 1.1089232734491513e-05, "grad_to_weight": 0.0003935773207301374, "update_to_weight": 9.175466865629975e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.771068690895454, "grad_norm": 0.03723061513756908, "grad_rms": 2.7780405611590195e-05, "grad_to_weight": 0.0007333037514777679, "update_to_weight": 7.101224836692348e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09054845216404, "grad_norm": 0.044683479502522105, "grad_rms": 3.837240867610042e-05, "grad_to_weight": 0.0011145639365806716, "update_to_weight": 8.6996201884337e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201484543703547, "grad_norm": 0.12016364881532215, "grad_rms": 0.00018731810026387415, "grad_to_weight": 0.004260897990285078, "update_to_weight": 7.542151642832009e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41176948719838, "grad_norm": 0.01438952705062212, "grad_rms": 0.0012249178876371754, "grad_to_weight": 0.0017106421036050868, "update_to_weight": 7.142054130602211e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3161751292382011, "grad_norm": 0.16201147317436604, "grad_rms": 0.0037505067496211874, "grad_to_weight": 0.12309264137830797, "update_to_weight": 0.00025943984631045903, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45700}
{"step": 45750, "lr": 1.1766910287607711e-05, "window_steps": 50, "elapsed_seconds": 2738.848478991, "steps_per_second": 1.7343055070172828, "eta_seconds": 2450.5486390972105, "images_seen": 46848000, "images_per_second": 1775.9288391856976, "effective_batch_size": 1024, "loss": 0.3515346348285675, "edge_loss": 0.18273061513900757, "diag_loss": 0.16880397498607635, "jvp_rms": 1.8825781345367432, "grad_norm": 0.28562089800834656, "grad_norm_postclip": 0.28562089800834656, "clip_fraction": 0.0, "ema_updates": 45750, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3536694223682086, "updates_nfe_2": 12, "loss_nfe_5": 0.3424840569496155, "updates_nfe_5": 11, "loss_nfe_4": 0.34423864632844925, "updates_nfe_4": 12, "loss_nfe_1": 0.4026344201781533, "updates_nfe_1": 11, "loss_nfe_3": 0.25138281285762787, "updates_nfe_3": 4, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.499577372369174, "grad_norm": 0.0021881946526865004, "grad_rms": 6.371809095020118e-06, "grad_to_weight": 0.0001326212546723127, "update_to_weight": 4.274029579007861e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.254297245807894, "grad_norm": 0.004046834229268297, "grad_rms": 1.1783985997969963e-05, "grad_to_weight": 0.00024897011344566166, "update_to_weight": 3.9860266093469546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062946441180251, "grad_norm": 0.0007168036242299338, "grad_rms": 1.2832843426910702e-05, "grad_to_weight": 0.0001415783541377609, "update_to_weight": 1.9231201917586102e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.855874086780574, "grad_norm": 0.004019783060484727, "grad_rms": 1.2863322258614156e-05, "grad_to_weight": 0.00031268064958867215, "update_to_weight": 5.778773112761197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0481478854491035, "grad_norm": 0.014639997534406281, "grad_rms": 0.00033307147670728893, "grad_to_weight": 0.002900073030072067, "update_to_weight": 3.197119019344165e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64106238421666, "grad_norm": 0.04048716178971326, "grad_rms": 9.755970923301744e-05, "grad_to_weight": 0.0024329673704074095, "update_to_weight": 6.24408319878491e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56762331269297, "grad_norm": 0.01573619355999617, "grad_rms": 2.1710444099753937e-05, "grad_to_weight": 0.000640525677217869, "update_to_weight": 6.581376254821983e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82922874829735, "grad_norm": 0.01002966869023939, "grad_rms": 1.2175784135144856e-05, "grad_to_weight": 0.0003253298605724385, "update_to_weight": 4.9813192857674355e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62339762481345, "grad_norm": 0.006691130698386382, "grad_rms": 6.006642999818604e-06, "grad_to_weight": 0.00019325459537197667, "update_to_weight": 5.376010081066746e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.068243652754845, "grad_norm": 0.005632213197997951, "grad_rms": 4.5303372310266296e-06, "grad_to_weight": 0.00015194173348915535, "update_to_weight": 5.929610554609694e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668801017729738, "grad_norm": 0.0019128640420777456, "grad_rms": 3.4559549528262464e-06, "grad_to_weight": 7.754183272640406e-05, "update_to_weight": 3.585613337457181e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3303748275533, "grad_norm": 0.00577083041196308, "grad_rms": 4.641835624857075e-06, "grad_to_weight": 0.00015458806504411708, "update_to_weight": 6.493688041950365e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.46450510484679, "grad_norm": 0.01403117165083638, "grad_rms": 6.229354059661797e-06, "grad_to_weight": 0.00022108691508199938, "update_to_weight": 7.305597765668674e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77153993703682, "grad_norm": 0.02170323591683712, "grad_rms": 1.619432541272631e-05, "grad_to_weight": 0.00042746853736861046, "update_to_weight": 6.218516500633424e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09123049614856, "grad_norm": 0.025275723504666777, "grad_rms": 2.1705793790084018e-05, "grad_to_weight": 0.0006304551691696002, "update_to_weight": 9.052093846079079e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20133412769742, "grad_norm": 0.06032032652090063, "grad_rms": 9.403084112864368e-05, "grad_to_weight": 0.0021389174798527753, "update_to_weight": 8.181554593311481e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412062440010114, "grad_norm": 0.009865406015247675, "grad_rms": 0.0008397991299066226, "grad_to_weight": 0.001172768995190175, "update_to_weight": 8.09680233814616e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3165669329007463, "grad_norm": 0.13119106104698727, "grad_rms": 0.0030370254050905206, "grad_to_weight": 0.09964632846879919, "update_to_weight": 0.00030613290302188234, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45750}
{"step": 45800, "lr": 1.1725856341707137e-05, "window_steps": 50, "elapsed_seconds": 2766.314817721, "steps_per_second": 1.7351604268795517, "eta_seconds": 2420.525465505875, "images_seen": 46899200, "images_per_second": 1776.804277124661, "effective_batch_size": 1024, "loss": 0.3522524833679199, "edge_loss": 0.18341979384422302, "diag_loss": 0.1688326746225357, "jvp_rms": 1.799648404121399, "grad_norm": 0.2578543424606323, "grad_norm_postclip": 0.2578543424606323, "clip_fraction": 0.0, "ema_updates": 45800, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3030436635017395, "updates_nfe_3": 8, "loss_nfe_1": 0.4062635451555252, "updates_nfe_1": 10, "loss_nfe_4": 0.3621487498283386, "updates_nfe_4": 10, "loss_nfe_2": 0.34787117080254987, "updates_nfe_2": 11, "loss_nfe_5": 0.3343245752833106, "updates_nfe_5": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49874873185227, "grad_norm": 0.19474066948177565, "grad_rms": 0.0005670658080856117, "grad_to_weight": 0.011803359918186514, "update_to_weight": 4.82794153741955e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25590712945797, "grad_norm": 0.12172612686353483, "grad_rms": 0.0003544545917331448, "grad_to_weight": 0.007488116528603323, "update_to_weight": 4.264324260720323e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.062759957880018, "grad_norm": 0.006379112033195063, "grad_rms": 0.00011420442525337283, "grad_to_weight": 0.001260006811752192, "update_to_weight": 2.5861348672189826e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.857240743113088, "grad_norm": 0.06781650653333134, "grad_rms": 0.00021701309868360434, "grad_to_weight": 0.0052745770175966315, "update_to_weight": 5.890808647625873e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0476268290984985, "grad_norm": 0.023977719588837817, "grad_rms": 0.0005455120093263992, "grad_to_weight": 0.004750295614289739, "update_to_weight": 2.5945211862066563e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64129895202135, "grad_norm": 0.1255778204755329, "grad_rms": 0.00030259803626990223, "grad_to_weight": 0.007546154950859738, "update_to_weight": 7.535048323424329e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.568630705558128, "grad_norm": 0.020071597255767775, "grad_rms": 2.7691785091020887e-05, "grad_to_weight": 0.0008169603547025111, "update_to_weight": 8.566320438783062e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.829666277044037, "grad_norm": 0.008393817250348693, "grad_rms": 1.0189898596507027e-05, "grad_to_weight": 0.00027226429163778465, "update_to_weight": 6.642722087955246e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.622193591141794, "grad_norm": 0.0034876373053724643, "grad_rms": 3.130859812867464e-06, "grad_to_weight": 0.00010073415181482863, "update_to_weight": 7.45355206981038e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.068467584491536, "grad_norm": 0.0029517860860007812, "grad_rms": 2.3743040139512446e-06, "grad_to_weight": 7.963064778096546e-05, "update_to_weight": 8.082448420973524e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668556072283724, "grad_norm": 0.0011130712252813315, "grad_rms": 2.010976174595817e-06, "grad_to_weight": 4.5121052972043185e-05, "update_to_weight": 5.173007528570155e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33080973497293, "grad_norm": 0.0031310331860018425, "grad_rms": 2.518483536661254e-06, "grad_to_weight": 8.387262982588269e-05, "update_to_weight": 7.888380170745408e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.4644185498851, "grad_norm": 0.008374068149552938, "grad_rms": 3.717796112927794e-06, "grad_to_weight": 0.000131949024995331, "update_to_weight": 8.680265489788875e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.771749111218334, "grad_norm": 0.028177021141528612, "grad_rms": 2.1024876256963366e-05, "grad_to_weight": 0.0005549744027885524, "update_to_weight": 7.197939285485444e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.091499961221515, "grad_norm": 0.08443004161678916, "grad_rms": 7.250518754424069e-05, "grad_to_weight": 0.0021059337190789586, "update_to_weight": 8.899532999234952e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20172802114981, "grad_norm": 0.1942995027671956, "grad_rms": 0.00030288539087642603, "grad_to_weight": 0.006889631111309251, "update_to_weight": 7.418043697652858e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411604413318088, "grad_norm": 0.010646784873584404, "grad_rms": 0.0009063145155222182, "grad_to_weight": 0.0012657258176249177, "update_to_weight": 5.253265101597958e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3161694684439391, "grad_norm": 0.15073087522790155, "grad_rms": 0.0034893650050952076, "grad_to_weight": 0.11452239156261948, "update_to_weight": 0.00012537284391494257, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45800}
{"step": 45850, "lr": 1.1685275634908482e-05, "window_steps": 50, "elapsed_seconds": 2793.79486623, "steps_per_second": 1.735990017958864, "eta_seconds": 2390.5667412071134, "images_seen": 46950400, "images_per_second": 1777.6537783898768, "effective_batch_size": 1024, "loss": 0.37260523438453674, "edge_loss": 0.19248171150684357, "diag_loss": 0.18012359738349915, "jvp_rms": 2.275195360183716, "grad_norm": 0.25538939237594604, "grad_norm_postclip": 0.25538939237594604, "clip_fraction": 0.0, "ema_updates": 45850, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.40179061591625215, "updates_nfe_1": 10, "loss_nfe_3": 0.3490452537169823, "updates_nfe_3": 13, "loss_nfe_4": 0.3455784499645233, "updates_nfe_4": 10, "loss_nfe_5": 0.39481428401036694, "updates_nfe_5": 11, "loss_nfe_2": 0.3793382793664932, "updates_nfe_2": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.500393949487425, "grad_norm": 0.08444022509767946, "grad_rms": 0.00024588168772023036, "grad_to_weight": 0.00511746721661167, "update_to_weight": 6.457127468534187e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.255263195130638, "grad_norm": 0.02692077131459641, "grad_rms": 7.839065656097166e-05, "grad_to_weight": 0.0016561264491036164, "update_to_weight": 5.0702037084949576e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06294003681717, "grad_norm": 0.004002045834849446, "grad_rms": 7.16481137230812e-05, "grad_to_weight": 0.0007904588649573148, "update_to_weight": 2.7748028638710658e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.856860966089942, "grad_norm": 0.03993792087298841, "grad_rms": 0.0001278015103796009, "grad_to_weight": 0.003106350840872041, "update_to_weight": 8.042892128846508e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0480473813041415, "grad_norm": 0.03999395853850549, "grad_rms": 0.0009098940623783605, "grad_to_weight": 0.007922659103126966, "update_to_weight": 2.6887771359951656e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64114399131994, "grad_norm": 0.22081989811268443, "grad_rms": 0.0005320976848076217, "grad_to_weight": 0.013269514297085862, "update_to_weight": 7.110166029216512e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56817732177262, "grad_norm": 0.03131169324038306, "grad_rms": 4.319918684097231e-05, "grad_to_weight": 0.001274481734248729, "update_to_weight": 7.34271690003476e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.829586096901302, "grad_norm": 0.012353584205489852, "grad_rms": 1.4996963431878755e-05, "grad_to_weight": 0.0004007054835786951, "update_to_weight": 5.9542667398404316e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.624192651914285, "grad_norm": 0.0033118177559359444, "grad_rms": 2.9730262099296503e-06, "grad_to_weight": 9.565039650831663e-05, "update_to_weight": 6.197928272601679e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06936659093496, "grad_norm": 0.002685478633587959, "grad_rms": 2.1600964681512176e-06, "grad_to_weight": 7.24446862883455e-05, "update_to_weight": 6.627298612605707e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66871689532552, "grad_norm": 0.0011752232922040927, "grad_rms": 2.12326577740355e-06, "grad_to_weight": 4.76402277909633e-05, "update_to_weight": 4.43634810855483e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33223704590394, "grad_norm": 0.0028302618009333836, "grad_rms": 2.27655451943452e-06, "grad_to_weight": 7.581281018475472e-05, "update_to_weight": 6.769150360344347e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.466590082031544, "grad_norm": 0.009050301128488735, "grad_rms": 4.018020125393624e-06, "grad_to_weight": 0.00014259945456012498, "update_to_weight": 7.426635096032734e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77175632410578, "grad_norm": 0.031674837009795086, "grad_rms": 2.36348450478643e-05, "grad_to_weight": 0.0006238672699757735, "update_to_weight": 6.482166721786677e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09185163325351, "grad_norm": 0.07566730877810654, "grad_rms": 6.498009842072165e-05, "grad_to_weight": 0.0018873488176671183, "update_to_weight": 8.443161878380197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20126595327216, "grad_norm": 0.19919178350108152, "grad_rms": 0.0003105117632616179, "grad_to_weight": 0.007063221340174253, "update_to_weight": 7.366732828345345e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41211050869217, "grad_norm": 0.009392912972769576, "grad_rms": 0.0007995778510918674, "grad_to_weight": 0.001116594101214428, "update_to_weight": 6.810135057927333e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3168013797525662, "grad_norm": 0.18127066354692853, "grad_rms": 0.004196350010405534, "grad_to_weight": 0.13765983718895422, "update_to_weight": 0.00010667686785917932, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45850}
{"step": 45900, "lr": 1.1645168610995847e-05, "window_steps": 50, "elapsed_seconds": 2821.376966488, "steps_per_second": 1.7367406263684901, "eta_seconds": 2360.74399236751, "images_seen": 47001600, "images_per_second": 1778.422401401334, "effective_batch_size": 1024, "loss": 0.3640049993991852, "edge_loss": 0.19186435639858246, "diag_loss": 0.17214062809944153, "jvp_rms": 2.060312509536743, "grad_norm": 0.23944351077079773, "grad_norm_postclip": 0.23944351077079773, "clip_fraction": 0.0, "ema_updates": 45900, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3589969327052434, "updates_nfe_3": 12, "loss_nfe_4": 0.3522354747567858, "updates_nfe_4": 7, "loss_nfe_2": 0.3479566350579262, "updates_nfe_2": 12, "loss_nfe_1": 0.4045160391501018, "updates_nfe_1": 14, "loss_nfe_5": 0.31758725047111513, "updates_nfe_5": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49912051544576, "grad_norm": 0.03301872208561783, "grad_rms": 9.614729358413412e-05, "grad_to_weight": 0.002001241342210158, "update_to_weight": 9.13838078524861e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25600850454994, "grad_norm": 0.016296836761089607, "grad_rms": 4.745479684217381e-05, "grad_to_weight": 0.0010025115794278922, "update_to_weight": 7.432974885815595e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063339352197106, "grad_norm": 0.0034085444766435496, "grad_rms": 6.1022734963735006e-05, "grad_to_weight": 0.0006731811240667682, "update_to_weight": 3.9707538047071264e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.858817884537261, "grad_norm": 0.020758886441270975, "grad_rms": 6.642852164062923e-05, "grad_to_weight": 0.0016143697366018033, "update_to_weight": 0.00011157688390114842, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.047188291480535, "grad_norm": 0.09597776379607813, "grad_rms": 0.0021835697337717853, "grad_to_weight": 0.019016085442678056, "update_to_weight": 6.311964848693698e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64153459955316, "grad_norm": 0.19008959209191173, "grad_rms": 0.0004580485396588512, "grad_to_weight": 0.011422599938410475, "update_to_weight": 0.000140619571607414, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56836364488893, "grad_norm": 0.05631819778096402, "grad_rms": 7.769941822721221e-05, "grad_to_weight": 0.0022923056087490045, "update_to_weight": 0.0001642122415126751, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.829665287167003, "grad_norm": 0.031322291323351605, "grad_rms": 3.802453197107024e-05, "grad_to_weight": 0.001015978961548754, "update_to_weight": 0.00012945207656324704, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.622711877239986, "grad_norm": 0.022244938643745583, "grad_rms": 1.9969331195110513e-05, "grad_to_weight": 0.0006424955596377992, "update_to_weight": 0.00014121607062490233, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06929743736722, "grad_norm": 0.018463410619771484, "grad_rms": 1.4851262479235723e-05, "grad_to_weight": 0.0004980782452369784, "update_to_weight": 0.0001398878864007959, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668815862830115, "grad_norm": 0.006668957963772419, "grad_rms": 1.2048748786168356e-05, "grad_to_weight": 0.0002703396061187076, "update_to_weight": 9.56591073078458e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33276512076891, "grad_norm": 0.016293202135481, "grad_rms": 1.3105629643645286e-05, "grad_to_weight": 0.00043643169968186446, "update_to_weight": 0.00013514701546892113, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.465764947036604, "grad_norm": 0.02746998437469391, "grad_rms": 1.2195721279850882e-05, "grad_to_weight": 0.00043283153362475246, "update_to_weight": 0.0001382326478792382, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.771996753100765, "grad_norm": 0.04956281825906882, "grad_rms": 3.698233804095995e-05, "grad_to_weight": 0.0009761841453683205, "update_to_weight": 0.00013086829474574506, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09223527195213, "grad_norm": 0.05513731282205015, "grad_rms": 4.734974814999085e-05, "grad_to_weight": 0.0013752616297905275, "update_to_weight": 0.00015745897502875754, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20169014701845, "grad_norm": 0.1557782801131707, "grad_rms": 0.00024283626355270932, "grad_to_weight": 0.005523721425952903, "update_to_weight": 0.0001348753011081232, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.411572667969715, "grad_norm": 0.014245202807380121, "grad_rms": 0.0012126321921765188, "grad_to_weight": 0.0016935243110511532, "update_to_weight": 9.577494492759104e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3162626646615396, "grad_norm": 0.1536128903218489, "grad_rms": 0.0035560826075623267, "grad_to_weight": 0.11670382701415338, "update_to_weight": 0.00026592639528325485, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45900}
{"step": 45950, "lr": 1.1605535708573227e-05, "window_steps": 50, "elapsed_seconds": 2848.926763928, "steps_per_second": 1.7374964013378547, "eta_seconds": 2330.9400795774545, "images_seen": 47052800, "images_per_second": 1779.1963149699632, "effective_batch_size": 1024, "loss": 0.3639390468597412, "edge_loss": 0.18813736736774445, "diag_loss": 0.17580164968967438, "jvp_rms": 2.136796712875366, "grad_norm": 0.24423758685588837, "grad_norm_postclip": 0.24423758685588837, "clip_fraction": 0.0, "ema_updates": 45950, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3323144001119277, "updates_nfe_3": 17, "loss_nfe_1": 0.40163793563842776, "updates_nfe_1": 10, "loss_nfe_2": 0.3333410455120934, "updates_nfe_2": 9, "loss_nfe_5": 0.3759945511817932, "updates_nfe_5": 5, "loss_nfe_4": 0.4056875771946377, "updates_nfe_4": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50105975822557, "grad_norm": 0.0923572389961975, "grad_rms": 0.0002689352588922669, "grad_to_weight": 0.005597048938033121, "update_to_weight": 5.1911692546625134e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25479665534579, "grad_norm": 0.025651718392336723, "grad_rms": 7.469529840707591e-05, "grad_to_weight": 0.0015781014636009319, "update_to_weight": 7.082345695299348e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063499257874644, "grad_norm": 0.0035847172678827826, "grad_rms": 6.417673386892147e-05, "grad_to_weight": 0.0007079525611281385, "update_to_weight": 3.1833858706248745e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.857457923056332, "grad_norm": 0.01951399801185239, "grad_rms": 6.244487356741696e-05, "grad_to_weight": 0.001517718209044991, "update_to_weight": 7.638418682691522e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.048092910685392, "grad_norm": 0.04430706747890764, "grad_rms": 0.001008020688465761, "grad_to_weight": 0.008776991284198048, "update_to_weight": 2.5767728775565058e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641375973371595, "grad_norm": 0.09593225803406066, "grad_rms": 0.00023116273865973164, "grad_to_weight": 0.00576468305190418, "update_to_weight": 7.384314808719158e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.568014598428455, "grad_norm": 0.04108632329725934, "grad_rms": 5.6684758090239456e-05, "grad_to_weight": 0.0016723501662152024, "update_to_weight": 9.45399379294277e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.829681125195734, "grad_norm": 0.028103838640369277, "grad_rms": 3.411740539217272e-05, "grad_to_weight": 0.0009115838248940322, "update_to_weight": 8.203164123561127e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.624187363545985, "grad_norm": 0.01991714925206065, "grad_rms": 1.7879669449601903e-05, "grad_to_weight": 0.0005752380277675595, "update_to_weight": 9.15709198215429e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.06989347102477, "grad_norm": 0.016576256819330185, "grad_rms": 1.3333308022921527e-05, "grad_to_weight": 0.00044716224588794176, "update_to_weight": 9.583898853314537e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669006374158673, "grad_norm": 0.006549548596624952, "grad_rms": 1.1833012913294249e-05, "grad_to_weight": 0.0002654970572096389, "update_to_weight": 6.670638480642793e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33408282299667, "grad_norm": 0.01583806902407091, "grad_rms": 1.2739537948034926e-05, "grad_to_weight": 0.00042422547512845657, "update_to_weight": 9.490326531108865e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.467495987119264, "grad_norm": 0.02918538695710687, "grad_rms": 1.2957300591017781e-05, "grad_to_weight": 0.00045984777724695564, "update_to_weight": 0.00010355993289286226, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.772364607259526, "grad_norm": 0.036213821059188016, "grad_rms": 2.7021703349580422e-05, "grad_to_weight": 0.000713258508625972, "update_to_weight": 7.677514743738356e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09262499639507, "grad_norm": 0.04144767430018031, "grad_rms": 3.5593626875692445e-05, "grad_to_weight": 0.0010337979691753055, "update_to_weight": 8.90548836006548e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201496447094236, "grad_norm": 0.10472884073593106, "grad_rms": 0.00016325742171530153, "grad_to_weight": 0.0037135916149840294, "update_to_weight": 7.585923428821255e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41241297230477, "grad_norm": 0.011706995781264672, "grad_rms": 0.0009965656614366672, "grad_to_weight": 0.0013916335086979542, "update_to_weight": 5.86679859566966e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3170873312007991, "grad_norm": 0.14875052724472337, "grad_rms": 0.003443520668691237, "grad_to_weight": 0.11293900087028118, "update_to_weight": 0.0001958814488442748, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 45950}
{"step": 46000, "lr": 1.1566377361059693e-05, "window_steps": 50, "elapsed_seconds": 2876.438703145, "steps_per_second": 1.7382605770577244, "eta_seconds": 2301.1509625159997, "images_seen": 47104000, "images_per_second": 1779.9788309071098, "effective_batch_size": 1024, "loss": 0.36331018805503845, "edge_loss": 0.18745647370815277, "diag_loss": 0.17585374414920807, "jvp_rms": 2.161757707595825, "grad_norm": 0.22839675843715668, "grad_norm_postclip": 0.22839675843715668, "clip_fraction": 0.0, "ema_updates": 46000, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3810098171234131, "updates_nfe_2": 11, "loss_nfe_1": 0.40494454652071, "updates_nfe_1": 8, "loss_nfe_3": 0.29508817331357434, "updates_nfe_3": 11, "loss_nfe_5": 0.386824294924736, "updates_nfe_5": 10, "loss_nfe_4": 0.3620632901787758, "updates_nfe_4": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.500163684379864, "grad_norm": 0.01075950832759545, "grad_rms": 3.1330637306670876e-05, "grad_to_weight": 0.0006520849449379164, "update_to_weight": 2.9697983230881008e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25522658585299, "grad_norm": 0.007640637934577877, "grad_rms": 2.2248791360278513e-05, "grad_to_weight": 0.0004700419212382782, "update_to_weight": 3.137865994662956e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063432207345507, "grad_norm": 0.0023912516794064, "grad_rms": 4.281027237986742e-05, "grad_to_weight": 0.0004722590490966618, "update_to_weight": 1.735539320343375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.857873877678879, "grad_norm": 0.008193846298509997, "grad_rms": 2.6220341717290867e-05, "grad_to_weight": 0.0006372629235953559, "update_to_weight": 4.1967394006358444e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.047590931285411, "grad_norm": 0.018542797269243225, "grad_rms": 0.00042186324514303914, "grad_to_weight": 0.0036735935066198296, "update_to_weight": 1.621026251920373e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64140714857868, "grad_norm": 0.04572255850337032, "grad_rms": 0.00011017515962582935, "grad_to_weight": 0.0027475175683851607, "update_to_weight": 5.096750913925448e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.568121424562182, "grad_norm": 0.02505817345893354, "grad_rms": 3.45715164247285e-05, "grad_to_weight": 0.001019946662827929, "update_to_weight": 6.071114989126553e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.829684094825215, "grad_norm": 0.01599077231409851, "grad_rms": 1.9412425062474036e-05, "grad_to_weight": 0.000518681030428806, "update_to_weight": 5.085323951799439e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62431252137903, "grad_norm": 0.012089606798512242, "grad_rms": 1.0852867074372847e-05, "grad_to_weight": 0.0003491652517590761, "update_to_weight": 5.8491499304379404e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07025898968741, "grad_norm": 0.0099926287946801, "grad_rms": 8.037689035006675e-06, "grad_to_weight": 0.00026955918482954094, "update_to_weight": 6.401136092113436e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66905833153839, "grad_norm": 0.003860945022586703, "grad_rms": 6.97553604432073e-06, "grad_to_weight": 0.00015650962313590387, "update_to_weight": 3.834375918349127e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.334785796476055, "grad_norm": 0.009645691918501502, "grad_rms": 7.758626259555157e-06, "grad_to_weight": 0.0002583566963818482, "update_to_weight": 6.822183350376863e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.4686846073026, "grad_norm": 0.020777522490909296, "grad_rms": 9.224500084477618e-06, "grad_to_weight": 0.00032736652129259145, "update_to_weight": 7.929713999981854e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77290796865736, "grad_norm": 0.02862565095444516, "grad_rms": 2.1359630816516617e-05, "grad_to_weight": 0.0005637977437123765, "update_to_weight": 6.455705108034327e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.092775709128695, "grad_norm": 0.02983506508226202, "grad_rms": 2.5621176393615285e-05, "grad_to_weight": 0.0007441506494515144, "update_to_weight": 7.035379924169058e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201863285489445, "grad_norm": 0.056726027509827304, "grad_rms": 8.842783831396279e-05, "grad_to_weight": 0.0020114283561900054, "update_to_weight": 5.121749153135774e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412217982015987, "grad_norm": 0.006465101613278344, "grad_rms": 0.0005503459970322112, "grad_to_weight": 0.0007685371000965174, "update_to_weight": 3.081642195160122e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.31686239518052, "grad_norm": 0.08732628043259437, "grad_rms": 0.0020215716687500543, "grad_to_weight": 0.0663138994265406, "update_to_weight": 5.6208192547800744e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46000}
{"step": 46050, "lr": 1.1527693996684657e-05, "window_steps": 50, "elapsed_seconds": 2927.180568702, "steps_per_second": 1.7252095938308725, "eta_seconds": 2289.5768804698814, "images_seen": 47155200, "images_per_second": 1766.6146240828134, "effective_batch_size": 1024, "loss": 0.3747604787349701, "edge_loss": 0.19403566420078278, "diag_loss": 0.1807248294353485, "jvp_rms": 2.3479297161102295, "grad_norm": 0.209913969039917, "grad_norm_postclip": 0.209913969039917, "clip_fraction": 0.0, "ema_updates": 46050, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3527640998363495, "updates_nfe_2": 15, "loss_nfe_5": 0.47073218652180265, "updates_nfe_5": 7, "loss_nfe_4": 0.3045142171057788, "updates_nfe_4": 11, "loss_nfe_1": 0.396895095705986, "updates_nfe_1": 10, "loss_nfe_3": 0.4046900314944131, "updates_nfe_3": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.500707437781738, "grad_norm": 0.046117912239287844, "grad_rms": 0.00013429085583810448, "grad_to_weight": 0.0027949051525931225, "update_to_weight": 3.248109795941065e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25392173980877, "grad_norm": 0.01959663042861748, "grad_rms": 5.706347369211958e-05, "grad_to_weight": 0.0012056555176232831, "update_to_weight": 3.768230960031466e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0633510298144495, "grad_norm": 0.002829852375888637, "grad_rms": 5.066248444274203e-05, "grad_to_weight": 0.0005588892334791054, "update_to_weight": 2.290298735891902e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.85759084011908, "grad_norm": 0.011173115446551914, "grad_rms": 3.575401519413486e-05, "grad_to_weight": 0.0008689898119707498, "update_to_weight": 4.370724812946837e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.04759225384095, "grad_norm": 0.028566708842489796, "grad_rms": 0.0006499151298676166, "grad_to_weight": 0.005659472359470408, "update_to_weight": 1.6572500472978113e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641722565036368, "grad_norm": 0.06313476237970449, "grad_rms": 0.00015213239920968142, "grad_to_weight": 0.003793763664366587, "update_to_weight": 5.3293977680754795e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56841333081463, "grad_norm": 0.028943603173324513, "grad_rms": 3.9932050679483095e-05, "grad_to_weight": 0.0011780819047448358, "update_to_weight": 7.174608838378946e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8298850390889, "grad_norm": 0.01566475694114697, "grad_rms": 1.9016649994683442e-05, "grad_to_weight": 0.0005081029955605018, "update_to_weight": 6.609220800649055e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62363384319752, "grad_norm": 0.011135067993616026, "grad_rms": 9.995975453369574e-06, "grad_to_weight": 0.00032160310047305235, "update_to_weight": 7.956175132081293e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.070099281525906, "grad_norm": 0.009476187520325991, "grad_rms": 7.62228339416962e-06, "grad_to_weight": 0.00025562886811712705, "update_to_weight": 8.647861734180658e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669120185419185, "grad_norm": 0.0035369590913669315, "grad_rms": 6.390193459058476e-06, "grad_to_weight": 0.00014337597225933782, "update_to_weight": 5.630087209415704e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.334710595294226, "grad_norm": 0.00928929501819292, "grad_rms": 7.471954201923463e-06, "grad_to_weight": 0.0002488112233917723, "update_to_weight": 8.730634091776583e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.46875192331361, "grad_norm": 0.02136353421078916, "grad_rms": 9.48466898391699e-06, "grad_to_weight": 0.00033659924866021854, "update_to_weight": 9.449396713948452e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.772799777692235, "grad_norm": 0.02921578290703403, "grad_rms": 2.1799970170209844e-05, "grad_to_weight": 0.0005754219392067169, "update_to_weight": 6.998331986051578e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09291880958376, "grad_norm": 0.027548848858365407, "grad_rms": 2.3657864130514998e-05, "grad_to_weight": 0.0006871250504161391, "update_to_weight": 7.0172395933868e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20185462859114, "grad_norm": 0.051840563066815656, "grad_rms": 8.081209156031595e-05, "grad_to_weight": 0.00183819694660292, "update_to_weight": 5.270766059488641e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412340418306702, "grad_norm": 0.005523557161013961, "grad_rms": 0.0004701964106332387, "grad_to_weight": 0.0006566017168056763, "update_to_weight": 2.7706864316030942e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3170505384597961, "grad_norm": 0.06560345250557974, "grad_rms": 0.0015186960935527128, "grad_to_weight": 0.04981088469262437, "update_to_weight": 6.229952078151583e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46050}
{"step": 46100, "lr": 1.148948603848319e-05, "window_steps": 50, "elapsed_seconds": 2954.687140947, "steps_per_second": 1.726071071729581, "eta_seconds": 2259.4666371947646, "images_seen": 47206400, "images_per_second": 1767.496777451091, "effective_batch_size": 1024, "loss": 0.386020302772522, "edge_loss": 0.1997339278459549, "diag_loss": 0.18628646433353424, "jvp_rms": 2.605585813522339, "grad_norm": 0.1884956955909729, "grad_norm_postclip": 0.1884956955909729, "clip_fraction": 0.0, "ema_updates": 46100, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.3977005048231645, "updates_nfe_1": 11, "loss_nfe_5": 0.385548429356681, "updates_nfe_5": 9, "loss_nfe_3": 0.403199123011695, "updates_nfe_3": 9, "loss_nfe_4": 0.39556538313627243, "updates_nfe_4": 8, "loss_nfe_2": 0.35869722870680004, "updates_nfe_2": 13, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.499218546539844, "grad_norm": 0.005617030973265554, "grad_rms": 1.635624554631011e-05, "grad_to_weight": 0.000340442243214212, "update_to_weight": 3.2001375348211673e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.256014136480957, "grad_norm": 0.01211518664452884, "grad_rms": 3.527824014865411e-05, "grad_to_weight": 0.000745274120876933, "update_to_weight": 4.679905825635454e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06367855717516, "grad_norm": 0.001315199219738224, "grad_rms": 2.3545843089489968e-05, "grad_to_weight": 0.0002597319724954906, "update_to_weight": 2.494963826375076e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.858412047215495, "grad_norm": 0.006974468764055719, "grad_rms": 2.2318328612457207e-05, "grad_to_weight": 0.0005424051382430265, "update_to_weight": 5.980045964461474e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046944538027129, "grad_norm": 0.02754736662682659, "grad_rms": 0.0006267242914645969, "grad_to_weight": 0.00545822654068534, "update_to_weight": 2.7430903458778767e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641294367405926, "grad_norm": 0.06015291406002036, "grad_rms": 0.00014494720167579903, "grad_to_weight": 0.003614677604515994, "update_to_weight": 6.703528021202314e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.568665485338613, "grad_norm": 0.02527745523201215, "grad_rms": 3.48740485878191e-05, "grad_to_weight": 0.0010288493384834642, "update_to_weight": 8.205133428859562e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.830030549634657, "grad_norm": 0.012030149678905823, "grad_rms": 1.4604321451453984e-05, "grad_to_weight": 0.000390208814731401, "update_to_weight": 7.111974716186125e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62363208037972, "grad_norm": 0.008196495597427442, "grad_rms": 7.358012438047963e-06, "grad_to_weight": 0.0002367312469818028, "update_to_weight": 8.393108937511455e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.070746342297994, "grad_norm": 0.006931595699108821, "grad_rms": 5.575510898141994e-06, "grad_to_weight": 0.0001869829011562095, "update_to_weight": 8.858168786998105e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668893799460538, "grad_norm": 0.002550634808987587, "grad_rms": 4.608209892113917e-06, "grad_to_weight": 0.00010339477844942381, "update_to_weight": 5.9230613074271756e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.335346530947554, "grad_norm": 0.007098853357641025, "grad_rms": 5.710046571949962e-06, "grad_to_weight": 0.00019013760463577675, "update_to_weight": 9.276173060515336e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.46939238350443, "grad_norm": 0.016606541370864404, "grad_rms": 7.372728983710307e-06, "grad_to_weight": 0.0002616464526794558, "update_to_weight": 9.927964223166197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.773141179509224, "grad_norm": 0.02980612970065131, "grad_rms": 2.224046983204972e-05, "grad_to_weight": 0.0005870452173772601, "update_to_weight": 7.46181468820392e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09341661250713, "grad_norm": 0.029810528228202585, "grad_rms": 2.5600105111743842e-05, "grad_to_weight": 0.0007435267619198909, "update_to_weight": 8.602718385834774e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201650108595707, "grad_norm": 0.05698440281906788, "grad_rms": 8.883060880702945e-05, "grad_to_weight": 0.002020605269537024, "update_to_weight": 6.440705439606307e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412319105450862, "grad_norm": 0.007178892608617916, "grad_rms": 0.000611107921670171, "grad_to_weight": 0.0008533785414733336, "update_to_weight": 2.8959100408008176e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3170562407231219, "grad_norm": 0.0718939864954872, "grad_rms": 0.0016643196702390234, "grad_to_weight": 0.054586876605978676, "update_to_weight": 0.00010067752116855503, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46100}
{"step": 46150, "lr": 1.1451753904291405e-05, "window_steps": 50, "elapsed_seconds": 2982.217820666, "steps_per_second": 1.7269026978216777, "eta_seconds": 2229.4249727308934, "images_seen": 47257600, "images_per_second": 1768.348362569398, "effective_batch_size": 1024, "loss": 0.3790767192840576, "edge_loss": 0.1962856501340866, "diag_loss": 0.1827910989522934, "jvp_rms": 2.251093626022339, "grad_norm": 0.42328011989593506, "grad_norm_postclip": 0.3474595248699188, "clip_fraction": 0.11999999731779099, "ema_updates": 46150, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3291581372419993, "updates_nfe_3": 12, "loss_nfe_1": 0.3976663459431041, "updates_nfe_1": 11, "loss_nfe_5": 0.43538326025009155, "updates_nfe_5": 10, "loss_nfe_2": 0.3488008603453636, "updates_nfe_2": 8, "loss_nfe_4": 0.3872632947232988, "updates_nfe_4": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50120123912148, "grad_norm": 0.3551404925508129, "grad_rms": 0.0010341344256860293, "grad_to_weight": 0.02152209935533884, "update_to_weight": 9.804243623881555e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25343826305972, "grad_norm": 0.21610909555745464, "grad_rms": 0.000629288577640466, "grad_to_weight": 0.013296207981336496, "update_to_weight": 9.529757827881644e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063976308059111, "grad_norm": 0.009259872549357336, "grad_rms": 0.0001657783116076058, "grad_to_weight": 0.0018285773838674223, "update_to_weight": 4.3499463533701445e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.857704767936385, "grad_norm": 0.09392604590664433, "grad_rms": 0.00030056373162308455, "grad_to_weight": 0.007305039865347532, "update_to_weight": 8.672553823319238e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.047292402739798, "grad_norm": 0.0511671815787326, "grad_rms": 0.0011640936883579609, "grad_to_weight": 0.010137550491617598, "update_to_weight": 2.567873889967736e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641811504061653, "grad_norm": 0.20590033331837568, "grad_rms": 0.0004961468218951776, "grad_to_weight": 0.01237247118609191, "update_to_weight": 7.482956788220063e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56866796960676, "grad_norm": 0.05258145501047957, "grad_rms": 7.254401995859926e-05, "grad_to_weight": 0.0021401833862351297, "update_to_weight": 8.533438541365719e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.82998006646189, "grad_norm": 0.03363321295775614, "grad_rms": 4.0829936999167354e-05, "grad_to_weight": 0.0010909255499111958, "update_to_weight": 7.435316343946636e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62370788146379, "grad_norm": 0.022015065089648254, "grad_rms": 1.976297319123904e-05, "grad_to_weight": 0.0006358378820956459, "update_to_weight": 8.640242160438039e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07102294498359, "grad_norm": 0.01873356308551765, "grad_rms": 1.5068562806939679e-05, "grad_to_weight": 0.0005053424911775372, "update_to_weight": 9.107208389762736e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668766379127447, "grad_norm": 0.006352346949839106, "grad_rms": 1.1476730400306668e-05, "grad_to_weight": 0.00025750565927017355, "update_to_weight": 5.784075059478118e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.335622808100446, "grad_norm": 0.01670768161540876, "grad_rms": 1.3439021110445717e-05, "grad_to_weight": 0.00044749974310817744, "update_to_weight": 9.38203818584701e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47025593671318, "grad_norm": 0.035178171052298556, "grad_rms": 1.5617889090755343e-05, "grad_to_weight": 0.0005542465605838277, "update_to_weight": 0.00010526267822244904, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77345853688593, "grad_norm": 0.05151847166145641, "grad_rms": 3.8441589910779066e-05, "grad_to_weight": 0.0010146732790327693, "update_to_weight": 7.8380646921148e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09362364798875, "grad_norm": 0.08827339375825098, "grad_rms": 7.580570667794008e-05, "grad_to_weight": 0.0022016816073615014, "update_to_weight": 8.506768545338052e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201893584612584, "grad_norm": 0.25747587884959827, "grad_rms": 0.00040136840854426795, "grad_to_weight": 0.00912973726665224, "update_to_weight": 7.88952937387648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412616573620012, "grad_norm": 0.0094069405864898, "grad_rms": 0.0008007719608708975, "grad_to_weight": 0.0011181943815183203, "update_to_weight": 4.9679976092250675e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3173637643261291, "grad_norm": 0.09300117648441894, "grad_rms": 0.0021529434508142783, "grad_to_weight": 0.07059642826291933, "update_to_weight": 0.00010119406610480543, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46150}
{"step": 46200, "lr": 1.1414498006741886e-05, "window_steps": 50, "elapsed_seconds": 3009.736358334, "steps_per_second": 1.727726079927609, "eta_seconds": 2199.422723397923, "images_seen": 47308800, "images_per_second": 1769.1915058458717, "effective_batch_size": 1024, "loss": 0.3544013500213623, "edge_loss": 0.18320921063423157, "diag_loss": 0.1711921989917755, "jvp_rms": 1.8829296827316284, "grad_norm": 0.30061936378479004, "grad_norm_postclip": 0.300207257270813, "clip_fraction": 0.03999999910593033, "ema_updates": 46200, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.36931730061769485, "updates_nfe_2": 12, "loss_nfe_5": 0.2561542157615934, "updates_nfe_5": 7, "loss_nfe_3": 0.34400975933441746, "updates_nfe_3": 13, "loss_nfe_1": 0.3961488664150238, "updates_nfe_1": 10, "loss_nfe_4": 0.38269593194127083, "updates_nfe_4": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49948396742778, "grad_norm": 0.030548409146921623, "grad_rms": 8.895398359637525e-05, "grad_to_weight": 0.001851476640556052, "update_to_weight": 3.2902802460182394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.255675278491683, "grad_norm": 0.008144291111209748, "grad_rms": 2.3715380215917445e-05, "grad_to_weight": 0.0005010121678541203, "update_to_weight": 4.4298412197150214e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063739766091971, "grad_norm": 0.0009731989987521969, "grad_rms": 1.7423056960168294e-05, "grad_to_weight": 0.00019218977350869674, "update_to_weight": 1.607763995253398e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.857835308986028, "grad_norm": 0.009343481306097762, "grad_rms": 2.989917845048575e-05, "grad_to_weight": 0.000726676075837418, "update_to_weight": 5.530081002311157e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.047272752140216, "grad_norm": 0.015394310897144502, "grad_rms": 0.00035023269992722503, "grad_to_weight": 0.003050025558974753, "update_to_weight": 2.8659520577476506e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641665717162958, "grad_norm": 0.04334699839647707, "grad_rms": 0.00010445090178583068, "grad_to_weight": 0.0026047271428948516, "update_to_weight": 5.9954041838897434e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56818353256593, "grad_norm": 0.017145834041678635, "grad_rms": 2.3655255007273357e-05, "grad_to_weight": 0.0006978877383812797, "update_to_weight": 7.295479335746327e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83040471645594, "grad_norm": 0.010016460912635052, "grad_rms": 1.2159750200826377e-05, "grad_to_weight": 0.0003248890504278297, "update_to_weight": 5.425630143024706e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.625021152971875, "grad_norm": 0.006495380484924284, "grad_rms": 5.830917595188761e-06, "grad_to_weight": 0.000187592101568053, "update_to_weight": 5.7236206829697296e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07178853094105, "grad_norm": 0.00542756623015917, "grad_rms": 4.365727024519125e-06, "grad_to_weight": 0.00014640691601995912, "update_to_weight": 6.191946612892031e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669065754012276, "grad_norm": 0.001926933293932258, "grad_rms": 3.481373748704885e-06, "grad_to_weight": 7.811132019131506e-05, "update_to_weight": 3.96085038796754e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33706791681693, "grad_norm": 0.0055776907546200944, "grad_rms": 4.486481459506999e-06, "grad_to_weight": 0.0001493874871761919, "update_to_weight": 6.642310165907014e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47196954537294, "grad_norm": 0.01394669806051168, "grad_rms": 6.19185071240626e-06, "grad_to_weight": 0.00021973003453976453, "update_to_weight": 7.724286006348564e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7734753663884, "grad_norm": 0.023638792207542895, "grad_rms": 1.7638581400471515e-05, "grad_to_weight": 0.000465573649173355, "update_to_weight": 6.859090169574146e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.094001180526526, "grad_norm": 0.02633200722414456, "grad_rms": 2.261288856007446e-05, "grad_to_weight": 0.0006567567827811084, "update_to_weight": 8.741696994437894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201544060606746, "grad_norm": 0.061472777293131484, "grad_rms": 9.582734857020527e-05, "grad_to_weight": 0.0021797663688563626, "update_to_weight": 7.403721174411452e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412724947366286, "grad_norm": 0.008034155456056194, "grad_rms": 0.0006839127301098868, "grad_to_weight": 0.0009550003722125007, "update_to_weight": 9.228656660064798e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3176558814584856, "grad_norm": 0.08173910126907129, "grad_rms": 0.0018922305007848482, "grad_to_weight": 0.06203372399370008, "update_to_weight": 0.0002923022624029293, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46200}
{"step": 46250, "lr": 1.137771875325914e-05, "window_steps": 50, "elapsed_seconds": 3037.246815545, "steps_per_second": 1.728539140490611, "eta_seconds": 2169.4620111035715, "images_seen": 47360000, "images_per_second": 1770.0240798623856, "effective_batch_size": 1024, "loss": 0.3760535717010498, "edge_loss": 0.1960468590259552, "diag_loss": 0.18000666797161102, "jvp_rms": 2.146601438522339, "grad_norm": 0.27425798773765564, "grad_norm_postclip": 0.27425798773765564, "clip_fraction": 0.0, "ema_updates": 46250, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3681894391775131, "updates_nfe_2": 10, "loss_nfe_3": 0.35209545989831287, "updates_nfe_3": 12, "loss_nfe_1": 0.40520108892367435, "updates_nfe_1": 13, "loss_nfe_5": 0.4098438819249471, "updates_nfe_5": 3, "loss_nfe_4": 0.36654109756151837, "updates_nfe_4": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.4994423511945, "grad_norm": 0.10202817347037389, "grad_rms": 0.0002970960754650735, "grad_to_weight": 0.00618373465591626, "update_to_weight": 5.39742368552358e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.255984099492984, "grad_norm": 0.04751278690246745, "grad_rms": 0.0001383525946118233, "grad_to_weight": 0.002922787486237105, "update_to_weight": 9.284801242354031e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063403578759212, "grad_norm": 0.004788922465663782, "grad_rms": 8.573546520708574e-05, "grad_to_weight": 0.0009457911839682564, "update_to_weight": 2.888367872651614e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.858667180360236, "grad_norm": 0.02348318961302836, "grad_rms": 7.514630294902008e-05, "grad_to_weight": 0.0018262537853764155, "update_to_weight": 8.390774587886472e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046953608133945, "grad_norm": 0.07524793195734278, "grad_rms": 0.0017119497293151815, "grad_to_weight": 0.014909574725646203, "update_to_weight": 3.336067436394108e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642046227454085, "grad_norm": 0.10512002495997776, "grad_rms": 0.00025330200034591254, "grad_to_weight": 0.006316532445785611, "update_to_weight": 7.31636546352355e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56887913148143, "grad_norm": 0.027033696692669563, "grad_rms": 3.7297047638504536e-05, "grad_to_weight": 0.0011003227517217026, "update_to_weight": 9.777153126628304e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.829969177923594, "grad_norm": 0.014381927764186536, "grad_rms": 1.7459325259108052e-05, "grad_to_weight": 0.000466491798327356, "update_to_weight": 7.673116547598266e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.6235033944384, "grad_norm": 0.010177386036864114, "grad_rms": 9.13626221791242e-06, "grad_to_weight": 0.00029394443193460644, "update_to_weight": 8.723470706025176e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07128966704726, "grad_norm": 0.008224106887767063, "grad_rms": 6.615157543900777e-06, "grad_to_weight": 0.00022184571838830541, "update_to_weight": 8.80380274576965e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668492979880988, "grad_norm": 0.003488305090683466, "grad_rms": 6.3022906960088895e-06, "grad_to_weight": 0.00014140730418872534, "update_to_weight": 6.0660469841829815e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33739812611646, "grad_norm": 0.008420660590885346, "grad_rms": 6.7732578373075385e-06, "grad_to_weight": 0.00022552885346864408, "update_to_weight": 8.842877474598395e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47140604055991, "grad_norm": 0.020967009410704854, "grad_rms": 9.30862571150684e-06, "grad_to_weight": 0.00033033787525214706, "update_to_weight": 9.798429345819957e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77355710960677, "grad_norm": 0.037191214478632065, "grad_rms": 2.7751005982210186e-05, "grad_to_weight": 0.0007324918047074386, "update_to_weight": 8.15217134054078e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.094428944952874, "grad_norm": 0.05687491286580255, "grad_rms": 4.884193048979081e-05, "grad_to_weight": 0.0014185240783423608, "update_to_weight": 0.00010918229991988804, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.201824329426145, "grad_norm": 0.17051087165876064, "grad_rms": 0.0002658022860353059, "grad_to_weight": 0.006046093673480811, "update_to_weight": 9.900057549633866e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412453330195351, "grad_norm": 0.026917767873484225, "grad_rms": 0.002291392569575148, "grad_to_weight": 0.00319975241667809, "update_to_weight": 1.1365603433806255e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3173726323938346, "grad_norm": 0.3346421306644097, "grad_rms": 0.0077468437584892365, "grad_to_weight": 0.25402237942071265, "update_to_weight": 0.0003950743114731114, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46250}
{"step": 46300, "lr": 1.1341416546055207e-05, "window_steps": 50, "elapsed_seconds": 3064.76917345, "steps_per_second": 1.7293308892277224, "eta_seconds": 2139.5558380688676, "images_seen": 47411200, "images_per_second": 1770.8348305691877, "effective_batch_size": 1024, "loss": 0.3683214485645294, "edge_loss": 0.19071798026561737, "diag_loss": 0.17760340869426727, "jvp_rms": 1.9969921112060547, "grad_norm": 0.30917105078697205, "grad_norm_postclip": 0.30917105078697205, "clip_fraction": 0.0, "ema_updates": 46300, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4070170289940304, "updates_nfe_1": 9, "loss_nfe_4": 0.4009522646665573, "updates_nfe_4": 12, "loss_nfe_5": 0.3352685086429119, "updates_nfe_5": 12, "loss_nfe_2": 0.3691052875735543, "updates_nfe_2": 11, "loss_nfe_3": 0.30968502163887024, "updates_nfe_3": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.500892383834596, "grad_norm": 0.028807795642853646, "grad_rms": 8.388548708829784e-05, "grad_to_weight": 0.0017458325872772636, "update_to_weight": 3.5704886177708264e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.253478631571554, "grad_norm": 0.009866930356218292, "grad_rms": 2.8731537437263417e-05, "grad_to_weight": 0.0006070657598830741, "update_to_weight": 6.969194446735246e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063815475790125, "grad_norm": 0.0015976900585599279, "grad_rms": 2.8603240375992423e-05, "grad_to_weight": 0.0003155111133488992, "update_to_weight": 2.527986926641343e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.857528535417003, "grad_norm": 0.010497271986019603, "grad_rms": 3.3591313352171336e-05, "grad_to_weight": 0.0008164299971883477, "update_to_weight": 6.20656568463365e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.047505909130762, "grad_norm": 0.042236475585944684, "grad_rms": 0.000960913091772079, "grad_to_weight": 0.008367791211406088, "update_to_weight": 2.050282493380577e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641842678452857, "grad_norm": 0.05881865966684752, "grad_rms": 0.00014173212151490983, "grad_to_weight": 0.0035343838301634344, "update_to_weight": 5.6859898472931544e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56903191217681, "grad_norm": 0.022699105338540495, "grad_rms": 3.1316827394625005e-05, "grad_to_weight": 0.0009238909135565268, "update_to_weight": 7.07109704450394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83015329269028, "grad_norm": 0.014135674058918014, "grad_rms": 1.7160378997727853e-05, "grad_to_weight": 0.00045850158202974395, "update_to_weight": 5.9290939780268e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62408688439487, "grad_norm": 0.010177373525858237, "grad_rms": 9.136250986754448e-06, "grad_to_weight": 0.00029393911700369477, "update_to_weight": 7.039347622422385e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07205360110396, "grad_norm": 0.008350245098819669, "grad_rms": 6.716618304297745e-06, "grad_to_weight": 0.00022524366167217153, "update_to_weight": 7.283196082691305e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668309887288842, "grad_norm": 0.0030826161805216023, "grad_rms": 5.569336043958616e-06, "grad_to_weight": 0.0001249626015972023, "update_to_weight": 4.935256666896009e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3382873881238, "grad_norm": 0.008455941432051705, "grad_rms": 6.80163639874671e-06, "grad_to_weight": 0.0002264683793381827, "update_to_weight": 7.454838122897671e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47309269362688, "grad_norm": 0.020030318238336262, "grad_rms": 8.892767285535988e-06, "grad_to_weight": 0.0003155718019762954, "update_to_weight": 8.6450376993767e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77386965599283, "grad_norm": 0.029209236353096933, "grad_rms": 2.179508532146215e-05, "grad_to_weight": 0.0005752808787472312, "update_to_weight": 6.536685203332756e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09477906902281, "grad_norm": 0.03581323122845421, "grad_rms": 3.075498953997885e-05, "grad_to_weight": 0.0008932143301451304, "update_to_weight": 7.551594306315449e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20191630893358, "grad_norm": 0.12327308146655425, "grad_rms": 0.00019216526513336326, "grad_to_weight": 0.004371088833686978, "update_to_weight": 6.021386807395916e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412790243008057, "grad_norm": 0.0202323156790032, "grad_rms": 0.0017222890854124101, "grad_to_weight": 0.002404947121535385, "update_to_weight": 4.69150165230196e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3176925215947288, "grad_norm": 0.23804682785148892, "grad_rms": 0.005510697588818563, "grad_to_weight": 0.1806543058796405, "update_to_weight": 0.00013654344401692418, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46300}
{"step": 46350, "lr": 1.1305591782125195e-05, "window_steps": 50, "elapsed_seconds": 3092.300327583, "steps_per_second": 1.7301036229497346, "eta_seconds": 2109.7002234912056, "images_seen": 47462400, "images_per_second": 1771.6261099005283, "effective_batch_size": 1024, "loss": 0.3764091432094574, "edge_loss": 0.19797717034816742, "diag_loss": 0.17843200266361237, "jvp_rms": 2.311210870742798, "grad_norm": 0.2329644411802292, "grad_norm_postclip": 0.2329644411802292, "clip_fraction": 0.0, "ema_updates": 46350, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4051567018032074, "updates_nfe_1": 16, "loss_nfe_5": 0.40528160482645037, "updates_nfe_5": 10, "loss_nfe_4": 0.37405784021724353, "updates_nfe_4": 11, "loss_nfe_2": 0.37261352439721424, "updates_nfe_2": 6, "loss_nfe_3": 0.2764027203832354, "updates_nfe_3": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.49802826168744, "grad_norm": 0.08061956616968234, "grad_rms": 0.00023475630210771482, "grad_to_weight": 0.004886618260734902, "update_to_weight": 5.812014550787298e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.257732722954195, "grad_norm": 0.0271425197432835, "grad_rms": 7.903636632585941e-05, "grad_to_weight": 0.001669514452341878, "update_to_weight": 9.15759698117006e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063789486021337, "grad_norm": 0.0031281437802153683, "grad_rms": 5.60027572286541e-05, "grad_to_weight": 0.000617747595718711, "update_to_weight": 2.409700592927534e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.858878403040434, "grad_norm": 0.018990361467787733, "grad_rms": 6.076923448159066e-05, "grad_to_weight": 0.0014768287616202618, "update_to_weight": 7.651999892557272e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046259886869205, "grad_norm": 0.030081571653981735, "grad_rms": 0.0006843794521768783, "grad_to_weight": 0.005961161796731185, "update_to_weight": 3.899363798479208e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641645545290228, "grad_norm": 0.057889207328395496, "grad_rms": 0.00013949247082375258, "grad_to_weight": 0.003478574710106044, "update_to_weight": 7.354643242480216e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.569056754476502, "grad_norm": 0.02519405365541836, "grad_rms": 3.475898357800162e-05, "grad_to_weight": 0.0010254383758883207, "update_to_weight": 8.64673237584006e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.830135475180267, "grad_norm": 0.011850538640739323, "grad_rms": 1.4386277835403854e-05, "grad_to_weight": 0.0003843816596355723, "update_to_weight": 6.427732586849677e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62430723302904, "grad_norm": 0.008603986047717883, "grad_rms": 7.723817527062332e-06, "grad_to_weight": 0.000248495543602106, "update_to_weight": 6.955108508149214e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07259526032589, "grad_norm": 0.007256985822968662, "grad_rms": 5.8372422887886765e-06, "grad_to_weight": 0.0001957506824653006, "update_to_weight": 7.195902100413579e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668303701704566, "grad_norm": 0.003126907430073829, "grad_rms": 5.64935666219901e-06, "grad_to_weight": 0.00012675810497086434, "update_to_weight": 4.753923095573178e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.338818646680394, "grad_norm": 0.007412062804700281, "grad_rms": 5.9619802912960605e-06, "grad_to_weight": 0.00019850823013007268, "update_to_weight": 7.652259987810579e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47349848374221, "grad_norm": 0.018489836465950833, "grad_rms": 8.208846753349202e-06, "grad_to_weight": 0.00029130010016206574, "update_to_weight": 8.595953179378182e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.773845614031444, "grad_norm": 0.028999686897551474, "grad_rms": 2.1638725593071133e-05, "grad_to_weight": 0.0005711540370213234, "update_to_weight": 7.907215495890078e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.094871927502304, "grad_norm": 0.031472710777941885, "grad_rms": 2.7027521884183845e-05, "grad_to_weight": 0.0007849560122014952, "update_to_weight": 0.00010884799595132239, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202120826998545, "grad_norm": 0.060444204510728405, "grad_rms": 9.42239493601914e-05, "grad_to_weight": 0.0021432503208362883, "update_to_weight": 9.207222033906325e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41225970124853, "grad_norm": 0.006901721722390169, "grad_rms": 0.000587513568966415, "grad_to_weight": 0.0008204361215055962, "update_to_weight": 5.561144471856778e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3172296501465797, "grad_norm": 0.07386120201564289, "grad_rms": 0.00170985999489469, "grad_to_weight": 0.05607313956789821, "update_to_weight": 0.00014346112728524623, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46350}
{"step": 46400, "lr": 1.1270244853243e-05, "window_steps": 50, "elapsed_seconds": 3119.836513062, "steps_per_second": 1.73085992724026, "eta_seconds": 2079.891008708, "images_seen": 47513600, "images_per_second": 1772.4005654940263, "effective_batch_size": 1024, "loss": 0.3713093400001526, "edge_loss": 0.19169433414936066, "diag_loss": 0.17961502075195312, "jvp_rms": 2.2922656536102295, "grad_norm": 0.17266994714736938, "grad_norm_postclip": 0.17266994714736938, "clip_fraction": 0.0, "ema_updates": 46400, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.37276743758808484, "updates_nfe_2": 11, "loss_nfe_5": 0.37798486029108364, "updates_nfe_5": 12, "loss_nfe_4": 0.31071288883686066, "updates_nfe_4": 4, "loss_nfe_1": 0.3971131377749973, "updates_nfe_1": 9, "loss_nfe_3": 0.36516691957201275, "updates_nfe_3": 14, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50057057636805, "grad_norm": 0.009567263023704732, "grad_rms": 2.7858935435222114e-05, "grad_to_weight": 0.0005798140724543713, "update_to_weight": 4.6283793913394895e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.252929423453175, "grad_norm": 0.004942047783445961, "grad_rms": 1.4390760427059907e-05, "grad_to_weight": 0.00030407120185451156, "update_to_weight": 3.703115309396786e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064267073809699, "grad_norm": 0.0007429037610863062, "grad_rms": 1.3300110832343528e-05, "grad_to_weight": 0.00014669521774005525, "update_to_weight": 2.7028103443567086e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.85765967162941, "grad_norm": 0.005870532486768992, "grad_rms": 1.8785728003408005e-05, "grad_to_weight": 0.00045657861824748686, "update_to_weight": 6.313583353058545e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0470161536351075, "grad_norm": 0.02018598305171438, "grad_rms": 0.00045924701612973863, "grad_to_weight": 0.003999587565650142, "update_to_weight": 2.627363695497478e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641850930487813, "grad_norm": 0.06630772757210268, "grad_rms": 0.0001597781206654023, "grad_to_weight": 0.003984396197818787, "update_to_weight": 7.223281280191873e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56915860764263, "grad_norm": 0.020202041743987872, "grad_rms": 2.7871752867778717e-05, "grad_to_weight": 0.0008222520789826194, "update_to_weight": 9.360259973896836e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.830261187392047, "grad_norm": 0.009825080321060563, "grad_rms": 1.192741862112674e-05, "grad_to_weight": 0.0003186830063275138, "update_to_weight": 8.042376663204897e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62464744856617, "grad_norm": 0.006427447087612586, "grad_rms": 5.769933632416279e-06, "grad_to_weight": 0.00018563213090213717, "update_to_weight": 9.693708575754726e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.072812580342315, "grad_norm": 0.00531540490537936, "grad_rms": 4.275508737734124e-06, "grad_to_weight": 0.0001433774384897311, "update_to_weight": 9.485988907125139e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668801017729738, "grad_norm": 0.0017738919735810286, "grad_rms": 3.204875316291254e-06, "grad_to_weight": 7.190831740489183e-05, "update_to_weight": 6.254352347488322e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33951989640132, "grad_norm": 0.005551755873382854, "grad_rms": 4.465620431360357e-06, "grad_to_weight": 0.0001486831081060021, "update_to_weight": 9.45960609196901e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.474363904232035, "grad_norm": 0.014302300748616131, "grad_rms": 6.349725984970406e-06, "grad_to_weight": 0.00022532405004002808, "update_to_weight": 9.803449473713989e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77435530117542, "grad_norm": 0.024108287948705672, "grad_rms": 1.7988905510729226e-05, "grad_to_weight": 0.00047481229068697934, "update_to_weight": 7.37610353372244e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.095491485951385, "grad_norm": 0.02815242033337069, "grad_rms": 2.4176187492124118e-05, "grad_to_weight": 0.0007021343121142362, "update_to_weight": 7.873066420741323e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202383776617072, "grad_norm": 0.0646149417446357, "grad_rms": 0.00010072553767793646, "grad_to_weight": 0.0022911163203944733, "update_to_weight": 5.947698259668484e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41286460687171, "grad_norm": 0.008693030944203662, "grad_rms": 0.0007399999363341259, "grad_to_weight": 0.001033302133152495, "update_to_weight": 2.2903600785690494e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.317750736604737, "grad_norm": 0.11752550335329794, "grad_rms": 0.002720672708807348, "grad_to_weight": 0.08918644481740862, "update_to_weight": 4.726366613181542e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46400}
{"step": 46450, "lr": 1.123537614595696e-05, "window_steps": 50, "elapsed_seconds": 3147.331859212, "steps_per_second": 1.731625466837336, "eta_seconds": 2050.096899119743, "images_seen": 47564800, "images_per_second": 1773.1844780414322, "effective_batch_size": 1024, "loss": 0.35307830572128296, "edge_loss": 0.1849178969860077, "diag_loss": 0.16816040873527527, "jvp_rms": 1.800624966621399, "grad_norm": 0.2310972660779953, "grad_norm_postclip": 0.2310972660779953, "clip_fraction": 0.0, "ema_updates": 46450, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3387933454730294, "updates_nfe_2": 11, "loss_nfe_5": 0.34522996470332146, "updates_nfe_5": 8, "loss_nfe_3": 0.3065229770210054, "updates_nfe_3": 9, "loss_nfe_4": 0.3545621534188588, "updates_nfe_4": 9, "loss_nfe_1": 0.4011986874617063, "updates_nfe_1": 13, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.500393024734386, "grad_norm": 0.07269644868568531, "grad_rms": 0.00021168495789093594, "grad_to_weight": 0.004405740431558934, "update_to_weight": 6.364116151894917e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.257487758049365, "grad_norm": 0.08288605118302067, "grad_rms": 0.00024135608508588162, "grad_to_weight": 0.005098330837861609, "update_to_weight": 7.015185112973323e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064232423845159, "grad_norm": 0.004131589891089744, "grad_rms": 7.396732436100816e-05, "grad_to_weight": 0.0008158373363031234, "update_to_weight": 3.369128854743783e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.85866777368723, "grad_norm": 0.032296738295203506, "grad_rms": 0.00010334969483234525, "grad_to_weight": 0.002511670638329463, "update_to_weight": 7.44135208431827e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046578507822911, "grad_norm": 0.04732936942881584, "grad_rms": 0.0010767804386737428, "grad_to_weight": 0.009378506517920729, "update_to_weight": 2.578541072466403e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64183992777363, "grad_norm": 0.14833261523738198, "grad_rms": 0.00035742902620575435, "grad_to_weight": 0.008913234106394035, "update_to_weight": 8.155652112616054e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56918469192214, "grad_norm": 0.03848139997512785, "grad_rms": 5.3090874858334494e-05, "grad_to_weight": 0.001566246518053152, "update_to_weight": 9.077014164067751e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.830124586696858, "grad_norm": 0.020901160419298494, "grad_rms": 2.5373521828001483e-05, "grad_to_weight": 0.0006779460251781567, "update_to_weight": 6.796624200502163e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.624721484665045, "grad_norm": 0.013796076673231364, "grad_rms": 1.2384768899254978e-05, "grad_to_weight": 0.0003984458526068293, "update_to_weight": 7.687288893056373e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07301014288815, "grad_norm": 0.011000140765691339, "grad_rms": 8.84809319275406e-06, "grad_to_weight": 0.00029671560855981736, "update_to_weight": 8.292597241856993e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668752771091814, "grad_norm": 0.004478786413971752, "grad_rms": 8.091784752879723e-06, "grad_to_weight": 0.00018155706757985095, "update_to_weight": 5.562701460062571e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.33979777726273, "grad_norm": 0.011151670313050137, "grad_rms": 8.96997777451034e-06, "grad_to_weight": 0.0002986537415004617, "update_to_weight": 8.518594506177584e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47505430900591, "grad_norm": 0.02248218945936617, "grad_rms": 9.981313155016543e-06, "grad_to_weight": 0.000354189369416205, "update_to_weight": 9.616262863661667e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.774098054058285, "grad_norm": 0.03419062745900394, "grad_rms": 2.5512054942316503e-05, "grad_to_weight": 0.0006733871948370561, "update_to_weight": 7.82347860266294e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09567111044377, "grad_norm": 0.0472428200814862, "grad_rms": 4.057026935594731e-05, "grad_to_weight": 0.0011782523841876987, "update_to_weight": 8.611213466733752e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20237403778603, "grad_norm": 0.12983298242131427, "grad_rms": 0.00020239122112653828, "grad_to_weight": 0.004603618909789714, "update_to_weight": 7.286966168434415e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412588459700448, "grad_norm": 0.006621403016158006, "grad_rms": 0.0005636512560290178, "grad_to_weight": 0.0007870827210765257, "update_to_weight": 7.380929896154025e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3176135857133564, "grad_norm": 0.07875182105221933, "grad_rms": 0.0018230760489624457, "grad_to_weight": 0.05976852539022893, "update_to_weight": 0.00016223810310346765, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46450}
{"step": 46500, "lr": 1.1200986041585675e-05, "window_steps": 50, "elapsed_seconds": 3174.809964919, "steps_per_second": 1.7323871541206795, "eta_seconds": 2020.3336140393635, "images_seen": 47616000, "images_per_second": 1773.9644458195758, "effective_batch_size": 1024, "loss": 0.31900879740715027, "edge_loss": 0.16345949470996857, "diag_loss": 0.1555492877960205, "jvp_rms": 1.298828125, "grad_norm": 0.16250088810920715, "grad_norm_postclip": 0.16250088810920715, "clip_fraction": 0.0, "ema_updates": 46500, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.3949831922849019, "updates_nfe_1": 6, "loss_nfe_4": 0.2567385733127594, "updates_nfe_4": 11, "loss_nfe_3": 0.3621643543243408, "updates_nfe_3": 10, "loss_nfe_2": 0.3547918266720242, "updates_nfe_2": 9, "loss_nfe_5": 0.2815460828798158, "updates_nfe_5": 14, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.502162907183845, "grad_norm": 0.006229052268178644, "grad_rms": 1.813839177744429e-05, "grad_to_weight": 0.0003774688386736848, "update_to_weight": 3.374411602818687e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25501631558461, "grad_norm": 0.0037182484790985204, "grad_rms": 1.0827176388342537e-05, "grad_to_weight": 0.00022874467837559928, "update_to_weight": 4.0886421006048906e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064502461314839, "grad_norm": 0.0006185042537365061, "grad_rms": 1.1073002393934288e-05, "grad_to_weight": 0.00012212537331375507, "update_to_weight": 2.3112938278173874e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.858851110417401, "grad_norm": 0.004803806251130701, "grad_rms": 1.5372199680046425e-05, "grad_to_weight": 0.0003735797397357662, "update_to_weight": 4.718735639148424e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046636333668373, "grad_norm": 0.00996961515127391, "grad_rms": 0.00022681659835216823, "grad_to_weight": 0.001975497042408651, "update_to_weight": 1.8179041533526357e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642065481960447, "grad_norm": 0.04924626366652626, "grad_rms": 0.00011866604009124443, "grad_to_weight": 0.0029591437264748113, "update_to_weight": 4.587741323635399e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.56962563626581, "grad_norm": 0.018156849001769615, "grad_rms": 2.5050102096017193e-05, "grad_to_weight": 0.0007389957531534114, "update_to_weight": 5.967990293598634e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.830573980907126, "grad_norm": 0.008914136625276558, "grad_rms": 1.0821554196120256e-05, "grad_to_weight": 0.00028913300903177924, "update_to_weight": 5.407854434208402e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.625442446970084, "grad_norm": 0.0059445756854404905, "grad_rms": 5.336458894227549e-06, "grad_to_weight": 0.00017168230253071246, "update_to_weight": 6.53700001037769e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07339703315741, "grad_norm": 0.004986827836788818, "grad_rms": 4.0112138904392e-06, "grad_to_weight": 0.0001345122981940052, "update_to_weight": 6.911596501064413e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668788646805933, "grad_norm": 0.001647899340497286, "grad_rms": 2.9772455136773556e-06, "grad_to_weight": 6.680098338394304e-05, "update_to_weight": 4.62550489346565e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.340239113204, "grad_norm": 0.0050969294919706405, "grad_rms": 4.0997754576478366e-06, "grad_to_weight": 0.0001364996479138319, "update_to_weight": 7.10859637312528e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47625048000236, "grad_norm": 0.013281045348505037, "grad_rms": 5.896323971870925e-06, "grad_to_weight": 0.00020922857364880294, "update_to_weight": 8.179753183731512e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.774939511750034, "grad_norm": 0.022377062282777138, "grad_rms": 1.6697115111162104e-05, "grad_to_weight": 0.00044071076200098225, "update_to_weight": 6.40574650055977e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09605471259766, "grad_norm": 0.02467754285441428, "grad_rms": 2.1192099856013178e-05, "grad_to_weight": 0.000615460624026456, "update_to_weight": 7.15160245675668e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202337246616192, "grad_norm": 0.04114403520172992, "grad_rms": 6.413772041012859e-05, "grad_to_weight": 0.001458887426313098, "update_to_weight": 4.958483927783636e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412891359564322, "grad_norm": 0.0037774854282539624, "grad_rms": 0.0003215609140647194, "grad_to_weight": 0.0004490115546255654, "update_to_weight": 3.73003728974784e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3178324232614917, "grad_norm": 0.05054150437856796, "grad_rms": 0.0011700174660088218, "grad_to_weight": 0.038351996419607924, "update_to_weight": 8.245291469754187e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46500}
{"step": 46550, "lr": 1.1167074916213815e-05, "window_steps": 50, "elapsed_seconds": 3204.323513169, "steps_per_second": 1.7320348514096138, "eta_seconds": 1991.8767784564052, "images_seen": 47667200, "images_per_second": 1773.6036878434445, "effective_batch_size": 1024, "loss": 0.35056743025779724, "edge_loss": 0.18252399563789368, "diag_loss": 0.1680433601140976, "jvp_rms": 1.9362499713897705, "grad_norm": 0.2238827496767044, "grad_norm_postclip": 0.21970342099666595, "clip_fraction": 0.019999999552965164, "ema_updates": 46550, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3779056340456009, "updates_nfe_2": 10, "loss_nfe_4": 0.3333374284780942, "updates_nfe_4": 13, "loss_nfe_3": 0.31067416444420815, "updates_nfe_3": 8, "loss_nfe_1": 0.40330918539654126, "updates_nfe_1": 11, "loss_nfe_5": 0.3117665443569422, "updates_nfe_5": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.501770849278856, "grad_norm": 0.009866206194566775, "grad_rms": 2.872942874926725e-05, "grad_to_weight": 0.0005978877227590357, "update_to_weight": 6.234112139244242e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25634547832159, "grad_norm": 0.017604324895713855, "grad_rms": 5.1262074580286105e-05, "grad_to_weight": 0.001082920199942222, "update_to_weight": 5.344531275973093e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064516019301906, "grad_norm": 0.0025498308712454247, "grad_rms": 4.564929532959437e-05, "grad_to_weight": 0.000503469800772176, "update_to_weight": 3.028540358888463e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.860152192112537, "grad_norm": 0.01173863363213111, "grad_rms": 3.756367570435522e-05, "grad_to_weight": 0.000912791190708514, "update_to_weight": 7.043770206742838e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046188638506355, "grad_norm": 0.05744797906363538, "grad_rms": 0.0013069867788984184, "grad_to_weight": 0.011384429552486898, "update_to_weight": 4.116490606693289e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64212141159068, "grad_norm": 0.10346735872671499, "grad_rms": 0.00024931966051152814, "grad_to_weight": 0.006217197685786227, "update_to_weight": 9.732229257807623e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570028068703444, "grad_norm": 0.034699299919366725, "grad_rms": 4.7872899397673076e-05, "grad_to_weight": 0.0014122613056175397, "update_to_weight": 0.00010746329326119054, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.830804614599238, "grad_norm": 0.018034641267915003, "grad_rms": 2.1893634357693533e-05, "grad_to_weight": 0.0005849552580076066, "update_to_weight": 8.829740633351394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62461219322535, "grad_norm": 0.012975981946170334, "grad_rms": 1.1648568027752463e-05, "grad_to_weight": 0.0003747617987389102, "update_to_weight": 0.00010591667870408487, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07349416667604, "grad_norm": 0.01087664483357074, "grad_rms": 8.748757780634867e-06, "grad_to_weight": 0.00029338062348995805, "update_to_weight": 0.00010642168821906557, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66865132913507, "grad_norm": 0.0039579285730008615, "grad_rms": 7.150755387684216e-06, "grad_to_weight": 0.00016044365458789083, "update_to_weight": 7.109679986389426e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34123291726252, "grad_norm": 0.010463704864686561, "grad_rms": 8.416604637731832e-06, "grad_to_weight": 0.0002802185157590038, "update_to_weight": 0.0001037918499357886, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47677548041193, "grad_norm": 0.02284648657527804, "grad_rms": 1.0143048452281857e-05, "grad_to_weight": 0.0003599188270413666, "update_to_weight": 0.00011016862095417754, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77510058920366, "grad_norm": 0.03507805114326263, "grad_rms": 2.6174224766987517e-05, "grad_to_weight": 0.0006908514357669495, "update_to_weight": 8.927027870439045e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09653573233447, "grad_norm": 0.03615611201493253, "grad_rms": 3.104944202694747e-05, "grad_to_weight": 0.0009017265794804232, "update_to_weight": 0.00010265719769514861, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202609932969857, "grad_norm": 0.08330700192585427, "grad_rms": 0.0001298638106721684, "grad_to_weight": 0.002953875620868174, "update_to_weight": 8.655826994589672e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412775732909237, "grad_norm": 0.008339286824409615, "grad_rms": 0.000709887237114907, "grad_to_weight": 0.0009912646062568687, "update_to_weight": 3.4333662901622207e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3177436803768616, "grad_norm": 0.08185571668219149, "grad_rms": 0.0018949301052353705, "grad_to_weight": 0.062118087076525816, "update_to_weight": 0.00010274269923113636, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46550}
{"step": 46600, "lr": 1.1133643140688021e-05, "window_steps": 50, "elapsed_seconds": 3231.811269852, "steps_per_second": 1.7327744513547818, "eta_seconds": 1962.1711281244286, "images_seen": 47718400, "images_per_second": 1774.3610381872966, "effective_batch_size": 1024, "loss": 0.36426106095314026, "edge_loss": 0.18835656344890594, "diag_loss": 0.17590448260307312, "jvp_rms": 2.149609327316284, "grad_norm": 0.21579991281032562, "grad_norm_postclip": 0.21579991281032562, "clip_fraction": 0.0, "ema_updates": 46600, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3247365706733295, "updates_nfe_3": 14, "loss_nfe_5": 0.33407509670807767, "updates_nfe_5": 13, "loss_nfe_2": 0.3701801101366679, "updates_nfe_2": 6, "loss_nfe_1": 0.3978437700054862, "updates_nfe_1": 11, "loss_nfe_4": 0.45440030097961426, "updates_nfe_4": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.5035174695332, "grad_norm": 0.03690543572222421, "grad_rms": 0.00010746502405617516, "grad_to_weight": 0.00223621635753436, "update_to_weight": 4.1704822250951774e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25472437333605, "grad_norm": 0.040344686981799306, "grad_rms": 0.0001174797877925344, "grad_to_weight": 0.0024820283663486774, "update_to_weight": 5.617801947480783e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064470072643259, "grad_norm": 0.003104306144804469, "grad_rms": 5.5575995096662245e-05, "grad_to_weight": 0.0006129577429182561, "update_to_weight": 2.6620722018413274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.859309144158525, "grad_norm": 0.016129227536074077, "grad_rms": 5.1613594180879875e-05, "grad_to_weight": 0.0012542841419595971, "update_to_weight": 5.616426677685676e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046395766236749, "grad_norm": 0.023386920552795162, "grad_rms": 0.0005320708658487784, "grad_to_weight": 0.004634380979246006, "update_to_weight": 2.3095759826145292e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641861933194722, "grad_norm": 0.03457974740230854, "grad_rms": 8.332493444323265e-05, "grad_to_weight": 0.0020778773157187408, "update_to_weight": 5.952021818996887e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.569183449814222, "grad_norm": 0.02203085265536838, "grad_rms": 3.0394872382620035e-05, "grad_to_weight": 0.0008966863998703614, "update_to_weight": 6.206250942134532e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.830854106583008, "grad_norm": 0.013349185208179522, "grad_rms": 1.6205599855261337e-05, "grad_to_weight": 0.00043298136217800084, "update_to_weight": 4.701225390123436e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.625867261250654, "grad_norm": 0.009956227955729976, "grad_rms": 8.937728113621488e-06, "grad_to_weight": 0.0002875372876760218, "update_to_weight": 5.46923099376353e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07400946396781, "grad_norm": 0.008026471148582404, "grad_rms": 6.4561869019996504e-06, "grad_to_weight": 0.00021649860008757138, "update_to_weight": 5.7703819092644195e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669561817621275, "grad_norm": 0.0031687331397251364, "grad_rms": 5.724922810783151e-06, "grad_to_weight": 0.00012844707835311996, "update_to_weight": 3.790921927959854e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.341685677775885, "grad_norm": 0.008077668524270611, "grad_rms": 6.497368116037072e-06, "grad_to_weight": 0.0002163177258245222, "update_to_weight": 6.30613997479554e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47860429408479, "grad_norm": 0.019441647933272486, "grad_rms": 8.631418066389859e-06, "grad_to_weight": 0.0003062708789752667, "update_to_weight": 7.36555654105434e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77559583906918, "grad_norm": 0.028089122273316332, "grad_rms": 2.095928866990057e-05, "grad_to_weight": 0.0005532012339617532, "update_to_weight": 6.158177164318832e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0967747178428, "grad_norm": 0.035038588907662455, "grad_rms": 3.0089757287652884e-05, "grad_to_weight": 0.0008738505566651104, "update_to_weight": 7.704637972886586e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202245268481615, "grad_norm": 0.06114546344799276, "grad_rms": 9.531711266886643e-05, "grad_to_weight": 0.0021681062222491895, "update_to_weight": 6.179805334453958e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41330397506565, "grad_norm": 0.006078630918932568, "grad_rms": 0.0005174474267813402, "grad_to_weight": 0.0007225022341933314, "update_to_weight": 3.927166383688662e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3182963038298587, "grad_norm": 0.075733457566834, "grad_rms": 0.0017532020307652028, "grad_to_weight": 0.057447978384537954, "update_to_weight": 0.00010214438291492194, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46600}
{"step": 46650, "lr": 1.1100691080612824e-05, "window_steps": 50, "elapsed_seconds": 3259.296931661, "steps_per_second": 1.7335026904470012, "eta_seconds": 1932.5034904538675, "images_seen": 47769600, "images_per_second": 1775.1067550177293, "effective_batch_size": 1024, "loss": 0.36552175879478455, "edge_loss": 0.18948589265346527, "diag_loss": 0.17603591084480286, "jvp_rms": 1.9964452981948853, "grad_norm": 0.24230998754501343, "grad_norm_postclip": 0.24230998754501343, "clip_fraction": 0.0, "ema_updates": 46650, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3564395010471344, "updates_nfe_2": 11, "loss_nfe_4": 0.33612911059306216, "updates_nfe_4": 13, "loss_nfe_5": 0.329144224524498, "updates_nfe_5": 6, "loss_nfe_3": 0.3932063529888789, "updates_nfe_3": 12, "loss_nfe_1": 0.41152939572930336, "updates_nfe_1": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.501920646087402, "grad_norm": 0.03382761002296171, "grad_rms": 9.850269624892539e-05, "grad_to_weight": 0.002049919566846434, "update_to_weight": 4.966047259481441e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25953277329516, "grad_norm": 0.015493870419859819, "grad_rms": 4.51166372868703e-05, "grad_to_weight": 0.000952909941256561, "update_to_weight": 5.510456317000462e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064226209424271, "grad_norm": 0.002766356708165704, "grad_rms": 4.952572964040558e-05, "grad_to_weight": 0.0005462545695564808, "update_to_weight": 3.1412242170728454e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.861278740651201, "grad_norm": 0.019858587796171594, "grad_rms": 6.354756228868089e-05, "grad_to_weight": 0.0015440601355916263, "update_to_weight": 7.500411187921008e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0457675533695285, "grad_norm": 0.019515742911041016, "grad_rms": 0.00044399852494127125, "grad_to_weight": 0.003867745135823496, "update_to_weight": 3.3790844252329274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64213516474965, "grad_norm": 0.11808663187885964, "grad_rms": 0.0002845469270047715, "grad_to_weight": 0.007095641917930309, "update_to_weight": 8.45707251091963e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570061604442383, "grad_norm": 0.019357652614360676, "grad_rms": 2.6706791155321518e-05, "grad_to_weight": 0.000787855273869591, "update_to_weight": 9.943675562971683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83094319195358, "grad_norm": 0.008444624467937256, "grad_rms": 1.0251577375036144e-05, "grad_to_weight": 0.00027390094475413835, "update_to_weight": 7.821214023374118e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62508108632202, "grad_norm": 0.003025633785430714, "grad_rms": 2.7161182192502744e-06, "grad_to_weight": 8.73827205743623e-05, "update_to_weight": 8.984834877906123e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07316325313703, "grad_norm": 0.002518787654918667, "grad_rms": 2.0260166099863683e-06, "grad_to_weight": 6.794099650251813e-05, "update_to_weight": 9.013966665537234e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66919193572661, "grad_norm": 0.0009990341036337942, "grad_rms": 1.8049462912927788e-06, "grad_to_weight": 4.049723664385477e-05, "update_to_weight": 6.10009280647803e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34108907883847, "grad_norm": 0.0026119394198795835, "grad_rms": 2.1009443327310063e-06, "grad_to_weight": 6.994813178493481e-05, "update_to_weight": 8.637968507920366e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.47847929798236, "grad_norm": 0.007320019640825788, "grad_rms": 3.24983509582139e-06, "grad_to_weight": 0.0001153149811050759, "update_to_weight": 9.168044853612292e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77580259243816, "grad_norm": 0.020355346856615185, "grad_rms": 1.5188569674497715e-05, "grad_to_weight": 0.00040088675741871237, "update_to_weight": 7.252622987919493e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09717809779449, "grad_norm": 0.042988695316409004, "grad_rms": 3.6916994905029763e-05, "grad_to_weight": 0.0010721127360025756, "update_to_weight": 9.513856197823181e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20262291796855, "grad_norm": 0.1040489922883942, "grad_rms": 0.00016219763432605814, "grad_to_weight": 0.00368933742762281, "update_to_weight": 7.804998513490117e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.412652849506426, "grad_norm": 0.008566365413192271, "grad_rms": 0.0007292174502845787, "grad_to_weight": 0.001018271592378243, "update_to_weight": 2.8042750586111393e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3177899522596033, "grad_norm": 0.13913900649657107, "grad_rms": 0.0032210174549757967, "grad_to_weight": 0.10558511715618304, "update_to_weight": 0.00010087330712452608, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46650}
{"step": 46700, "lr": 1.1068219096346685e-05, "window_steps": 50, "elapsed_seconds": 3286.7728991520003, "steps_per_second": 1.7342238648342942, "eta_seconds": 1902.8685205616844, "images_seen": 47820800, "images_per_second": 1775.8452375903173, "effective_batch_size": 1024, "loss": 0.3685298562049866, "edge_loss": 0.1899634301662445, "diag_loss": 0.17856645584106445, "jvp_rms": 2.0998828411102295, "grad_norm": 0.22773335874080658, "grad_norm_postclip": 0.22773335874080658, "clip_fraction": 0.0, "ema_updates": 46700, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3729828119277954, "updates_nfe_2": 10, "loss_nfe_3": 0.35676304357392447, "updates_nfe_3": 7, "loss_nfe_4": 0.39700796741705674, "updates_nfe_4": 13, "loss_nfe_1": 0.393210510412852, "updates_nfe_1": 9, "loss_nfe_5": 0.31812056221745233, "updates_nfe_5": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.505259282555212, "grad_norm": 0.020543331653873443, "grad_rms": 5.982017519029149e-05, "grad_to_weight": 0.001244653676878991, "update_to_weight": 5.282799798328531e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.255205934429227, "grad_norm": 0.0070486579044658805, "grad_rms": 2.0525003334699095e-05, "grad_to_weight": 0.00043362464510748026, "update_to_weight": 7.629671404584433e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0647613749350695, "grad_norm": 0.0014807936686852252, "grad_rms": 2.6510459288222744e-05, "grad_to_weight": 0.0002923718530972664, "update_to_weight": 2.6769658214394787e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.860105324606053, "grad_norm": 0.0080882116747217, "grad_rms": 2.5882310488488067e-05, "grad_to_weight": 0.0006289382140009393, "update_to_weight": 6.912529911152332e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046226058196355, "grad_norm": 0.015576381031712844, "grad_rms": 0.00035437493891616294, "grad_to_weight": 0.0030867386542092856, "update_to_weight": 1.7414209476833933e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.641853681165223, "grad_norm": 0.03739608103408659, "grad_rms": 9.01113002459658e-05, "grad_to_weight": 0.002247110312982166, "update_to_weight": 5.420527670761347e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.569139975997455, "grad_norm": 0.02138407075178755, "grad_rms": 2.9502539538029156e-05, "grad_to_weight": 0.0008703630152572893, "update_to_weight": 6.74707500685676e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831158975450975, "grad_norm": 0.01272064613397314, "grad_rms": 1.544256806184942e-05, "grad_to_weight": 0.0004125905920079695, "update_to_weight": 5.6869418149176926e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62721922885864, "grad_norm": 0.009455510357711073, "grad_rms": 8.488232805488956e-06, "grad_to_weight": 0.0002730658299535287, "update_to_weight": 6.858718192563401e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.073818492004904, "grad_norm": 0.008020692881241033, "grad_rms": 6.451539084392756e-06, "grad_to_weight": 0.00021634385686412967, "update_to_weight": 7.124039451816144e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669848812461083, "grad_norm": 0.0029549061217997836, "grad_rms": 5.338603383269361e-06, "grad_to_weight": 0.00011977803934928127, "update_to_weight": 4.827895862928491e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34213843279968, "grad_norm": 0.008161375637311664, "grad_rms": 6.564698921420369e-06, "grad_to_weight": 0.00021855672920282127, "update_to_weight": 7.488979716844354e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48023883584757, "grad_norm": 0.019059365958548864, "grad_rms": 8.461698114953147e-06, "grad_to_weight": 0.0003002409302182076, "update_to_weight": 8.551463509400799e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.776026173622284, "grad_norm": 0.02572797677157717, "grad_rms": 1.919747035172549e-05, "grad_to_weight": 0.0005066953582307438, "update_to_weight": 6.416931072428143e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09747644412925, "grad_norm": 0.025477036511425452, "grad_rms": 2.187867345507793e-05, "grad_to_weight": 0.0006353775541691377, "update_to_weight": 6.946254855445385e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202450866250636, "grad_norm": 0.057178162264505554, "grad_rms": 8.913265232505999e-05, "grad_to_weight": 0.002027418203321103, "update_to_weight": 5.3798339834259375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.413533398667191, "grad_norm": 0.007390045107398741, "grad_rms": 0.0006290824160274916, "grad_to_weight": 0.000878352144958433, "update_to_weight": 3.530726829510457e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3185516443275365, "grad_norm": 0.09247995938379341, "grad_rms": 0.0021408774642788053, "grad_to_weight": 0.07013753293748183, "update_to_weight": 0.00010594639351468035, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46700}
{"step": 46750, "lr": 1.1036227542998006e-05, "window_steps": 50, "elapsed_seconds": 3314.247236143, "steps_per_second": 1.7349339353124542, "eta_seconds": 1873.2701769503915, "images_seen": 47872000, "images_per_second": 1776.572349759953, "effective_batch_size": 1024, "loss": 0.3530244827270508, "edge_loss": 0.1830107718706131, "diag_loss": 0.17001362144947052, "jvp_rms": 1.8420312404632568, "grad_norm": 0.19888602197170258, "grad_norm_postclip": 0.19888602197170258, "clip_fraction": 0.0, "ema_updates": 46750, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.40352954864501955, "updates_nfe_1": 10, "loss_nfe_3": 0.3295769492785136, "updates_nfe_3": 9, "loss_nfe_5": 0.34758039850455064, "updates_nfe_5": 13, "loss_nfe_4": 0.3219121504913677, "updates_nfe_4": 11, "loss_nfe_2": 0.37002178175108774, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50495050320604, "grad_norm": 0.02238661842501766, "grad_rms": 6.518764622340424e-05, "grad_to_weight": 0.0013563578043247766, "update_to_weight": 6.113477353823569e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.256504107010525, "grad_norm": 0.010166271989303447, "grad_rms": 2.9603191034382155e-05, "grad_to_weight": 0.0006253664331754636, "update_to_weight": 4.368200909112322e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064106815736268, "grad_norm": 0.0030973143981139145, "grad_rms": 5.54508227516467e-05, "grad_to_weight": 0.0006116210638545936, "update_to_weight": 2.966174025542263e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.861159505585995, "grad_norm": 0.014044995266154072, "grad_rms": 4.4944042380104094e-05, "grad_to_weight": 0.0010920473585646691, "update_to_weight": 6.576200894563912e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045680232393682, "grad_norm": 0.03167021678924872, "grad_rms": 0.0007205223804747611, "grad_to_weight": 0.006276699142748549, "update_to_weight": 2.6577963949952865e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642285531879608, "grad_norm": 0.05574343902810917, "grad_rms": 0.0001343219297879384, "grad_to_weight": 0.0033495062274546497, "update_to_weight": 8.455190350016407e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570239218889128, "grad_norm": 0.025490302657945753, "grad_rms": 3.5167703601967844e-05, "grad_to_weight": 0.001037446254831304, "update_to_weight": 0.00010234185094669253, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831460871846268, "grad_norm": 0.014269787718796558, "grad_rms": 1.732318985645986e-05, "grad_to_weight": 0.0004628320331011952, "update_to_weight": 9.239553094344948e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62645952474423, "grad_norm": 0.010779650748155728, "grad_rms": 9.676916596848696e-06, "grad_to_weight": 0.0003113125308249473, "update_to_weight": 0.00010898158988540416, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.074017695495826, "grad_norm": 0.00870693474120221, "grad_rms": 7.003525832475409e-06, "grad_to_weight": 0.0002348527427676129, "update_to_weight": 0.00010974271401688256, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669052146141784, "grad_norm": 0.003527800883045097, "grad_rms": 6.373647403137805e-06, "grad_to_weight": 0.00014300512488870965, "update_to_weight": 7.450996813230324e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34178211342156, "grad_norm": 0.008863642210033607, "grad_rms": 7.129575336545882e-06, "grad_to_weight": 0.00023736527043918975, "update_to_weight": 0.00010669523331106896, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.479944621436346, "grad_norm": 0.02039798351903619, "grad_rms": 9.055997931266723e-06, "grad_to_weight": 0.0003213295733113803, "update_to_weight": 0.00011042308445894962, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77609108410424, "grad_norm": 0.03071155317757319, "grad_rms": 2.2916070580149217e-05, "grad_to_weight": 0.0006048428014417917, "update_to_weight": 8.205846746065336e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.097630182955484, "grad_norm": 0.038863687920807896, "grad_rms": 3.3374601355149476e-05, "grad_to_weight": 0.0009692265538756925, "update_to_weight": 8.139877721386741e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20282093845789, "grad_norm": 0.10359779731185961, "grad_rms": 0.00016149428529591215, "grad_to_weight": 0.0036733133021665835, "update_to_weight": 6.453603941611982e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41319152801193, "grad_norm": 0.01473638439104526, "grad_rms": 0.0012544443452648564, "grad_to_weight": 0.0017515807576684902, "update_to_weight": 2.8848206291675267e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.31818340116065, "grad_norm": 0.1936653298291089, "grad_rms": 0.004483282032192586, "grad_to_weight": 0.14691834964587486, "update_to_weight": 8.088684525954828e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46750}
{"step": 46800, "lr": 1.10047167704213e-05, "window_steps": 50, "elapsed_seconds": 3341.740100125, "steps_per_second": 1.7356227073981747, "eta_seconds": 1843.7186759310346, "images_seen": 47923200, "images_per_second": 1777.277652375731, "effective_batch_size": 1024, "loss": 0.36336612701416016, "edge_loss": 0.18815484642982483, "diag_loss": 0.17521123588085175, "jvp_rms": 2.0317578315734863, "grad_norm": 0.22432661056518555, "grad_norm_postclip": 0.22432661056518555, "clip_fraction": 0.0, "ema_updates": 46800, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.37516065858877623, "updates_nfe_5": 13, "loss_nfe_2": 0.3665254024358896, "updates_nfe_2": 13, "loss_nfe_3": 0.334106864200698, "updates_nfe_3": 9, "loss_nfe_1": 0.4007057934999466, "updates_nfe_1": 10, "loss_nfe_4": 0.30247327089309695, "updates_nfe_4": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.504909825256355, "grad_norm": 0.07895506432414598, "grad_rms": 0.00022990943531501644, "grad_to_weight": 0.004783731941590274, "update_to_weight": 5.059213884713059e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.254086024750823, "grad_norm": 0.023725759810198997, "grad_rms": 6.908709513538308e-05, "grad_to_weight": 0.0014596797244748627, "update_to_weight": 6.832891528860587e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.063890807699618, "grad_norm": 0.004161680867317391, "grad_rms": 7.450603925228277e-05, "grad_to_weight": 0.0008218346377037944, "update_to_weight": 3.07065465485617e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.860363390402432, "grad_norm": 0.015078421780341226, "grad_rms": 4.825101145842611e-05, "grad_to_weight": 0.0011724724506302917, "update_to_weight": 6.25877268586894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.046252516395559, "grad_norm": 0.02400932760218453, "grad_rms": 0.0005462311165295549, "grad_to_weight": 0.00475785298579032, "update_to_weight": 2.5461077743321786e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642526666574888, "grad_norm": 0.04703960970550974, "grad_rms": 0.00011334878619400149, "grad_to_weight": 0.0028264704421345366, "update_to_weight": 6.811086923330692e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.569815674627197, "grad_norm": 0.02177629147540433, "grad_rms": 3.004366697538442e-05, "grad_to_weight": 0.0008863025984314694, "update_to_weight": 8.375464809011333e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8314044520888, "grad_norm": 0.014859779974072453, "grad_rms": 1.803942670969086e-05, "grad_to_weight": 0.0004819689611339037, "update_to_weight": 6.755243674542848e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62640840716357, "grad_norm": 0.010920971276723354, "grad_rms": 9.803780351558509e-06, "grad_to_weight": 0.00031539428370122287, "update_to_weight": 7.872206027081088e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0743090904099, "grad_norm": 0.008948537465664588, "grad_rms": 7.197861838459423e-06, "grad_to_weight": 0.00024136761237660738, "update_to_weight": 8.251726724244632e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668949468331515, "grad_norm": 0.0038052132251502908, "grad_rms": 6.874845886973703e-06, "grad_to_weight": 0.0001542511256928549, "update_to_weight": 5.534158315698482e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.342551955154825, "grad_norm": 0.009127956804939443, "grad_rms": 7.342179903864335e-06, "grad_to_weight": 0.00024443848443730703, "update_to_weight": 8.528368356753878e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48064073437566, "grad_norm": 0.021270645913033535, "grad_rms": 9.443429798115667e-06, "grad_to_weight": 0.0003350729555808529, "update_to_weight": 9.501897051084087e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77611512500243, "grad_norm": 0.028842147788396, "grad_rms": 2.1521174477252707e-05, "grad_to_weight": 0.0005680258861354672, "update_to_weight": 7.062178258614991e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09823752166359, "grad_norm": 0.030339567765871194, "grad_rms": 2.6054423387116495e-05, "grad_to_weight": 0.0007566309553999687, "update_to_weight": 8.732876681446894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202797132789527, "grad_norm": 0.06429417474135828, "grad_rms": 0.00010022550737531469, "grad_to_weight": 0.0022797091522034772, "update_to_weight": 7.136071491711586e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.413498033381527, "grad_norm": 0.009127808940406533, "grad_rms": 0.0007770106972039154, "grad_to_weight": 0.0010849005852489528, "update_to_weight": 5.44891424131817e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.318616194965134, "grad_norm": 0.09855995225395335, "grad_rms": 0.002281627090526828, "grad_to_weight": 0.07474498844340328, "update_to_weight": 0.00017977501908979918, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46800}
{"step": 46850, "lr": 1.0973687123213323e-05, "window_steps": 50, "elapsed_seconds": 3369.260794482, "steps_per_second": 1.7362858967702426, "eta_seconds": 1814.2173508749233, "images_seen": 47974400, "images_per_second": 1777.9567582927284, "effective_batch_size": 1024, "loss": 0.37645453214645386, "edge_loss": 0.1955944448709488, "diag_loss": 0.18086005747318268, "jvp_rms": 2.301953077316284, "grad_norm": 0.21492986381053925, "grad_norm_postclip": 0.21492986381053925, "clip_fraction": 0.0, "ema_updates": 46850, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.40875188329003076, "updates_nfe_1": 11, "loss_nfe_3": 0.3577991894313267, "updates_nfe_3": 14, "loss_nfe_2": 0.376254881421725, "updates_nfe_2": 6, "loss_nfe_4": 0.39718289034707205, "updates_nfe_4": 14, "loss_nfe_5": 0.29983547925949094, "updates_nfe_5": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.5030893843703, "grad_norm": 0.06384656614048054, "grad_rms": 0.00018591496433841332, "grad_to_weight": 0.0038687644872691636, "update_to_weight": 6.082800025831986e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25688424638592, "grad_norm": 0.10077489770046816, "grad_rms": 0.0002934466588377869, "grad_to_weight": 0.006198906024865836, "update_to_weight": 6.574507297486972e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064476098458306, "grad_norm": 0.003960334438529146, "grad_rms": 7.09013599400332e-05, "grad_to_weight": 0.0007819830445511876, "update_to_weight": 2.639194590647451e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.861000523779488, "grad_norm": 0.0393333877165147, "grad_rms": 0.00012586700180271246, "grad_to_weight": 0.00305834586071191, "update_to_weight": 7.655234118725996e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045402004833793, "grad_norm": 0.02762518741592099, "grad_rms": 0.0006284947757205731, "grad_to_weight": 0.005475319387722609, "update_to_weight": 2.628391165899484e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64182892505217, "grad_norm": 0.12951629955030028, "grad_rms": 0.00031208837484562955, "grad_to_weight": 0.007782576069829072, "update_to_weight": 6.798931839662627e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.569409512421387, "grad_norm": 0.020021684170307436, "grad_rms": 2.7622922487901433e-05, "grad_to_weight": 0.0008149029450701781, "update_to_weight": 9.366980316089122e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831702386834117, "grad_norm": 0.008332882894757447, "grad_rms": 1.011592570836814e-05, "grad_to_weight": 0.00027026995753292525, "update_to_weight": 8.661604181229683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.627506537207886, "grad_norm": 0.0030369310514977957, "grad_rms": 2.7262598002771233e-06, "grad_to_weight": 8.770285114911632e-05, "update_to_weight": 9.968295330285942e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07529849898801, "grad_norm": 0.0026026494591714923, "grad_rms": 2.0934718430735483e-06, "grad_to_weight": 7.01990156395513e-05, "update_to_weight": 9.536263883065846e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668971735844732, "grad_norm": 0.0010081056476278652, "grad_rms": 1.8213357715205724e-06, "grad_to_weight": 4.08653290628672e-05, "update_to_weight": 6.532147375978752e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34370260195737, "grad_norm": 0.002635762325561202, "grad_rms": 2.120106568386261e-06, "grad_to_weight": 7.058117277913007e-05, "update_to_weight": 9.126082346079596e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48198678746889, "grad_norm": 0.0074019253749051436, "grad_rms": 3.2861984038753114e-06, "grad_to_weight": 0.00011659882983319382, "update_to_weight": 9.394828678805826e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77630985585822, "grad_norm": 0.02579853334676504, "grad_rms": 1.9250117622527783e-05, "grad_to_weight": 0.0005080820843421054, "update_to_weight": 7.032147536471145e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09843235519002, "grad_norm": 0.05044645434421419, "grad_rms": 4.332142402310382e-05, "grad_to_weight": 0.0012580654998520112, "update_to_weight": 8.219984620563718e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20262291796855, "grad_norm": 0.1619826763903345, "grad_rms": 0.0002525080381316316, "grad_to_weight": 0.005743532325396996, "update_to_weight": 6.31385085260485e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.413333900238035, "grad_norm": 0.011803771216356243, "grad_rms": 0.001004803733550539, "grad_to_weight": 0.0014029838059823447, "update_to_weight": 4.810454974381359e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.318482976700125, "grad_norm": 0.20105102761040863, "grad_rms": 0.004654258252806385, "grad_to_weight": 0.15248663135082371, "update_to_weight": 0.00010891128683124643, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46850}
{"step": 46900, "lr": 1.0943138940709337e-05, "window_steps": 50, "elapsed_seconds": 3396.766902527, "steps_per_second": 1.7369457985505976, "eta_seconds": 1784.7419318362204, "images_seen": 48025600, "images_per_second": 1778.632497715812, "effective_batch_size": 1024, "loss": 0.3632211983203888, "edge_loss": 0.18844959139823914, "diag_loss": 0.17477159202098846, "jvp_rms": 2.1788671016693115, "grad_norm": 0.17903664708137512, "grad_norm_postclip": 0.17903664708137512, "clip_fraction": 0.0, "ema_updates": 46900, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4054675069120195, "updates_nfe_1": 9, "loss_nfe_2": 0.37379874537388486, "updates_nfe_2": 12, "loss_nfe_5": 0.3433950424194336, "updates_nfe_5": 10, "loss_nfe_3": 0.3561919331550598, "updates_nfe_3": 14, "loss_nfe_4": 0.3211260139942169, "updates_nfe_4": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50600901914919, "grad_norm": 0.08327284184500602, "grad_rms": 0.00024248238171350626, "grad_to_weight": 0.005045001595988366, "update_to_weight": 7.543149883055396e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2560554372488, "grad_norm": 0.026868669559968843, "grad_rms": 7.823894134057532e-05, "grad_to_weight": 0.0016528406699698201, "update_to_weight": 8.009763230111707e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064005875028122, "grad_norm": 0.004011467519176588, "grad_rms": 7.181678893021047e-05, "grad_to_weight": 0.0007921530144659067, "update_to_weight": 2.896279953257357e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.861656607659814, "grad_norm": 0.020380691942094763, "grad_rms": 6.521829769417771e-05, "grad_to_weight": 0.0015846086210976084, "update_to_weight": 6.84459956047404e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045540931518243, "grad_norm": 0.06941524187974128, "grad_rms": 0.0015792514352917688, "grad_to_weight": 0.013757740314050667, "update_to_weight": 2.9517934102304887e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64248540803448, "grad_norm": 0.15492984254925554, "grad_rms": 0.00037332600563922097, "grad_to_weight": 0.009309297184329221, "update_to_weight": 6.727761657962927e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570283932813282, "grad_norm": 0.05127039863216008, "grad_rms": 7.073522063844537e-05, "grad_to_weight": 0.002086683197164407, "update_to_weight": 8.232278094492668e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83173604041175, "grad_norm": 0.03360624538399074, "grad_rms": 4.0797198992862555e-05, "grad_to_weight": 0.0010899887486044376, "update_to_weight": 7.399026659991943e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.626413695192646, "grad_norm": 0.024496017229644725, "grad_rms": 2.1990129478619373e-05, "grad_to_weight": 0.000707437317802439, "update_to_weight": 8.9372946940066e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07490668976373, "grad_norm": 0.019931750502837204, "grad_rms": 1.6032339012776478e-05, "grad_to_weight": 0.0005376075702529091, "update_to_weight": 9.659421913849301e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66876514203361, "grad_norm": 0.006991254958526251, "grad_rms": 1.2631039984496632e-05, "grad_to_weight": 0.0002834051448571988, "update_to_weight": 6.404012755797555e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.343735290269066, "grad_norm": 0.016603610845029572, "grad_rms": 1.3355310556683644e-05, "grad_to_weight": 0.00044461569567081033, "update_to_weight": 9.384288967332618e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.481792572998835, "grad_norm": 0.0289106158716225, "grad_rms": 1.2835311749356244e-05, "grad_to_weight": 0.00045541587122603183, "update_to_weight": 0.00010152685674847157, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77639159451344, "grad_norm": 0.03820180385544093, "grad_rms": 2.850507847579584e-05, "grad_to_weight": 0.0007523536560161704, "update_to_weight": 7.669974536930524e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09891486851998, "grad_norm": 0.0372677594674219, "grad_rms": 3.200408098581056e-05, "grad_to_weight": 0.0009293957103233059, "update_to_weight": 8.285959461501356e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203020039624636, "grad_norm": 0.08713137590439453, "grad_rms": 0.0001358254677574966, "grad_to_weight": 0.00308943424434606, "update_to_weight": 6.878612253506456e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.413444985174307, "grad_norm": 0.004281167754902015, "grad_rms": 0.000364437201063411, "grad_to_weight": 0.0005088483685869515, "update_to_weight": 5.385633500019439e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3184968099660404, "grad_norm": 0.04318976591739873, "grad_rms": 0.0009998273913195464, "grad_to_weight": 0.032756822459442385, "update_to_weight": 8.109031186868795e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46900}
{"step": 46950, "lr": 1.0913072556979356e-05, "window_steps": 50, "elapsed_seconds": 3424.254864502, "steps_per_second": 1.7376043067592537, "eta_seconds": 1755.2903086943024, "images_seen": 48076800, "images_per_second": 1779.3068101214758, "effective_batch_size": 1024, "loss": 0.3805314302444458, "edge_loss": 0.1951371282339096, "diag_loss": 0.1853943020105362, "jvp_rms": 2.4377734661102295, "grad_norm": 0.28503161668777466, "grad_norm_postclip": 0.27992159128189087, "clip_fraction": 0.03999999910593033, "ema_updates": 46950, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.3986393851893289, "updates_nfe_1": 7, "loss_nfe_2": 0.34798019751906395, "updates_nfe_2": 8, "loss_nfe_3": 0.37782662237683934, "updates_nfe_3": 12, "loss_nfe_4": 0.3722920740644137, "updates_nfe_4": 6, "loss_nfe_5": 0.3932107231196235, "updates_nfe_5": 17, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50420071852684, "grad_norm": 0.21959003021305606, "grad_rms": 0.0006394247193545961, "grad_to_weight": 0.013305099347619703, "update_to_weight": 6.52899923095012e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.255894926857977, "grad_norm": 0.22303457911015004, "grad_rms": 0.0006494549092939594, "grad_to_weight": 0.013720227653640433, "update_to_weight": 5.810679245943587e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064574016947973, "grad_norm": 0.010767681244477124, "grad_rms": 0.0001927724174521375, "grad_to_weight": 0.002126078364822867, "update_to_weight": 3.276023821740009e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.861257978405899, "grad_norm": 0.16124219938380088, "grad_rms": 0.0005159756984774795, "grad_to_weight": 0.012537047282196435, "update_to_weight": 7.320787034281628e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045609920899252, "grad_norm": 0.04668330443906029, "grad_rms": 0.0010620819512128644, "grad_to_weight": 0.009252261901122188, "update_to_weight": 2.846318268628102e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642164504784134, "grad_norm": 0.25484675302145315, "grad_rms": 0.0006140903443142455, "grad_to_weight": 0.015313317744707558, "update_to_weight": 7.43453591651092e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.569778412319245, "grad_norm": 0.052594150765654935, "grad_rms": 7.256153566098231e-05, "grad_to_weight": 0.002140603382051029, "update_to_weight": 6.624179519642213e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831897379112004, "grad_norm": 0.01847126757988673, "grad_rms": 2.242368852862201e-05, "grad_to_weight": 0.0005990960385201802, "update_to_weight": 4.523680542453394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62733203732909, "grad_norm": 0.004363120135950753, "grad_rms": 3.916782708832212e-06, "grad_to_weight": 0.00012600220343996486, "update_to_weight": 3.422107062471409e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07564585580405, "grad_norm": 0.0031997799986824647, "grad_rms": 2.5737808476920483e-06, "grad_to_weight": 8.630409328881728e-05, "update_to_weight": 3.8408325639461274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668877717318864, "grad_norm": 0.001307809860921576, "grad_rms": 2.3628087866075466e-06, "grad_to_weight": 5.3014566609303995e-05, "update_to_weight": 2.1914966847882413e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.344782935505194, "grad_norm": 0.0035745303648886846, "grad_rms": 2.8752157324667206e-06, "grad_to_weight": 9.571699401926993e-05, "update_to_weight": 4.4825803530325446e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.482880935532535, "grad_norm": 0.011863290834587868, "grad_rms": 5.266890090178849e-06, "grad_to_weight": 0.00018687385732596402, "update_to_weight": 5.1326978123129986e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77664161781847, "grad_norm": 0.04563678823911638, "grad_rms": 3.405284826501746e-05, "grad_to_weight": 0.0008987752396586539, "update_to_weight": 6.295150346666227e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09900467319918, "grad_norm": 0.15787995851209805, "grad_rms": 0.00013558107732971098, "grad_to_weight": 0.003937253799659014, "update_to_weight": 8.971603215399529e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.202618589636316, "grad_norm": 0.46086710052817914, "grad_rms": 0.0007184265008275164, "grad_to_weight": 0.01634128756744365, "update_to_weight": 7.943318432735859e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.413643120786906, "grad_norm": 0.007762440754037157, "grad_rms": 0.0006607828386501995, "grad_to_weight": 0.0009226016177057859, "update_to_weight": 4.504833408910591e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.31884743027754, "grad_norm": 0.0832927525576601, "grad_rms": 0.001928196963716388, "grad_to_weight": 0.06315571509293677, "update_to_weight": 0.00010537525344818926, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 46950}
{"step": 47000, "lr": 1.088348830082454e-05, "window_steps": 50, "elapsed_seconds": 3451.742206638, "steps_per_second": 1.7382526390474582, "eta_seconds": 1725.871103319, "images_seen": 48128000, "images_per_second": 1779.9707023845972, "effective_batch_size": 1024, "loss": 0.3659606873989105, "edge_loss": 0.1881421059370041, "diag_loss": 0.17781850695610046, "jvp_rms": 2.2069921493530273, "grad_norm": 0.22301813960075378, "grad_norm_postclip": 0.22301813960075378, "clip_fraction": 0.0, "ema_updates": 47000, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.41566345947129385, "updates_nfe_5": 7, "loss_nfe_1": 0.39811037853360176, "updates_nfe_1": 8, "loss_nfe_4": 0.36302856107552844, "updates_nfe_4": 12, "loss_nfe_2": 0.35789306089282036, "updates_nfe_2": 8, "loss_nfe_3": 0.33226783871650695, "updates_nfe_3": 15, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.505246339824232, "grad_norm": 0.024794378400512917, "grad_rms": 7.219880809223094e-05, "grad_to_weight": 0.0015022119567333254, "update_to_weight": 2.0359415576384567e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2552096892355, "grad_norm": 0.019119946923460557, "grad_rms": 5.56754178855322e-05, "grad_to_weight": 0.001176235021817168, "update_to_weight": 2.802020820633379e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0648986408602, "grad_norm": 0.0027784305700010196, "grad_rms": 4.9741886441589745e-05, "grad_to_weight": 0.0005485658780190601, "update_to_weight": 1.9803684635610872e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.861733722011877, "grad_norm": 0.012183410321526165, "grad_rms": 3.898696293222822e-05, "grad_to_weight": 0.0009472603449000959, "update_to_weight": 3.548240422182917e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045509366199244, "grad_norm": 0.02635107607806761, "grad_rms": 0.0005995077390908886, "grad_to_weight": 0.005222679052901597, "update_to_weight": 1.7910635365921872e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642516581163346, "grad_norm": 0.054085616746816736, "grad_rms": 0.00013032716570536443, "grad_to_weight": 0.0032498460484055016, "update_to_weight": 4.7338301272125736e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570307531796, "grad_norm": 0.02306740465716678, "grad_rms": 3.182495165850927e-05, "grad_to_weight": 0.0009388325574401403, "update_to_weight": 6.186543766881717e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831926083454874, "grad_norm": 0.014746286990150067, "grad_rms": 1.7901648871182872e-05, "grad_to_weight": 0.00047827978538334865, "update_to_weight": 5.3009541395998406e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.627841433501885, "grad_norm": 0.011135820716444528, "grad_rms": 9.996651174336992e-06, "grad_to_weight": 0.00032158576034343274, "update_to_weight": 6.205014935504896e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.076164415611245, "grad_norm": 0.009195868375941255, "grad_rms": 7.396805378382306e-06, "grad_to_weight": 0.0002480264213109718, "update_to_weight": 6.642242464677907e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66891235576482, "grad_norm": 0.0034796894228384173, "grad_rms": 6.286724843284375e-06, "grad_to_weight": 0.0001410556482043383, "update_to_weight": 4.430831446552958e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34474371062814, "grad_norm": 0.008966436388131991, "grad_rms": 7.2122590482238124e-06, "grad_to_weight": 0.00024009902056390832, "update_to_weight": 6.693256711001699e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.482888627074146, "grad_norm": 0.020644175217165023, "grad_rms": 9.16529851516361e-06, "grad_to_weight": 0.0003251927513638802, "update_to_weight": 7.566463791996444e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.777067135420495, "grad_norm": 0.026685310032425972, "grad_rms": 1.991180467560318e-05, "grad_to_weight": 0.0005255386247743941, "update_to_weight": 5.564486487943307e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.0995145773105, "grad_norm": 0.028546579384943527, "grad_rms": 2.4514675729359165e-05, "grad_to_weight": 0.0007118933903777487, "update_to_weight": 5.928376780712106e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203010301013304, "grad_norm": 0.048290251329504694, "grad_rms": 7.527765867204947e-05, "grad_to_weight": 0.0017122374815347167, "update_to_weight": 4.407869288603297e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.413717930502692, "grad_norm": 0.004977915745883204, "grad_rms": 0.00042374832882498845, "grad_to_weight": 0.0005916428132011064, "update_to_weight": 2.615935867013388e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.318867225318187, "grad_norm": 0.053405770438057325, "grad_rms": 0.0012363241847759803, "grad_to_weight": 0.040493667150742006, "update_to_weight": 6.316511854822516e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47000}
{"step": 47050, "lr": 1.0854386495773575e-05, "window_steps": 50, "elapsed_seconds": 3502.445186476, "steps_per_second": 1.7273646489489343, "eta_seconds": 1707.8038512568924, "images_seen": 48179200, "images_per_second": 1768.8214005237087, "effective_batch_size": 1024, "loss": 0.35724833607673645, "edge_loss": 0.18516509234905243, "diag_loss": 0.17208316922187805, "jvp_rms": 1.9724608659744263, "grad_norm": 0.1658777892589569, "grad_norm_postclip": 0.1658777892589569, "clip_fraction": 0.0, "ema_updates": 47050, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.3963787718252702, "updates_nfe_1": 11, "loss_nfe_3": 0.3189087788263957, "updates_nfe_3": 15, "loss_nfe_4": 0.41915935940212673, "updates_nfe_4": 9, "loss_nfe_2": 0.37232260406017303, "updates_nfe_2": 8, "loss_nfe_5": 0.2810861383165632, "updates_nfe_5": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50346569308436, "grad_norm": 0.0246401333460503, "grad_rms": 7.174966155964064e-05, "grad_to_weight": 0.0014930278163559029, "update_to_weight": 3.7609408214271046e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.259069171458297, "grad_norm": 0.04812611910676846, "grad_rms": 0.000140138558083051, "grad_to_weight": 0.0029599553700928112, "update_to_weight": 5.0118721949367784e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064931403370593, "grad_norm": 0.002831852501076696, "grad_rms": 5.069829242766978e-05, "grad_to_weight": 0.0005591097441501704, "update_to_weight": 2.2201402689864598e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.862637111870944, "grad_norm": 0.018078418230940245, "grad_rms": 5.785101238835203e-05, "grad_to_weight": 0.0014054985827327469, "update_to_weight": 4.820893745346192e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045040777550123, "grad_norm": 0.022733560253402976, "grad_rms": 0.0005172064043467159, "grad_to_weight": 0.004506120219000968, "update_to_weight": 2.159716105839824e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642105824663442, "grad_norm": 0.04539739265316132, "grad_rms": 0.00010939162518190701, "grad_to_weight": 0.0027278634766210183, "update_to_weight": 4.914960885614572e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570492596716747, "grad_norm": 0.02744608381900758, "grad_rms": 3.78659976593393e-05, "grad_to_weight": 0.0011170343333969253, "update_to_weight": 6.588922976107474e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83214680906065, "grad_norm": 0.014949797123906249, "grad_rms": 1.8148705432516864e-05, "grad_to_weight": 0.00048487694407036715, "update_to_weight": 5.5027426110704996e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62856938111669, "grad_norm": 0.010696002650385823, "grad_rms": 9.60182551231213e-06, "grad_to_weight": 0.0003088779825890948, "update_to_weight": 6.514708665548328e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07696940361024, "grad_norm": 0.008631730630157602, "grad_rms": 6.943034517211853e-06, "grad_to_weight": 0.00023280572196165302, "update_to_weight": 6.827373275174261e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.668901221983926, "grad_norm": 0.0034195094755106615, "grad_rms": 6.177998251925444e-06, "grad_to_weight": 0.0001386162052675185, "update_to_weight": 4.5712844810882456e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34566875300769, "grad_norm": 0.0086361443554175, "grad_rms": 6.9465847492733004e-06, "grad_to_weight": 0.00023124888758946038, "update_to_weight": 7.259272027648261e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48436730863847, "grad_norm": 0.020468407904017525, "grad_rms": 9.087263918127949e-06, "grad_to_weight": 0.0003224165061692651, "update_to_weight": 8.073458152486967e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77761044649748, "grad_norm": 0.031289965435577334, "grad_rms": 2.334766503752514e-05, "grad_to_weight": 0.0006162157919689118, "update_to_weight": 6.31380359954463e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09958154930906, "grad_norm": 0.03656310417468801, "grad_rms": 3.139895083100613e-05, "grad_to_weight": 0.0009118076239705305, "update_to_weight": 7.233913204153637e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20280578939853, "grad_norm": 0.07369737008513282, "grad_rms": 0.0001148837564012993, "grad_to_weight": 0.002613121922529983, "update_to_weight": 5.728401045495968e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41376689650194, "grad_norm": 0.008900190683426993, "grad_rms": 0.0007576345444265397, "grad_to_weight": 0.0010578128432732414, "update_to_weight": 4.851986477995986e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3190217339299106, "grad_norm": 0.10002468985253475, "grad_rms": 0.0023155352338346236, "grad_to_weight": 0.0758324804508868, "update_to_weight": 0.0001614798977243567, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47050}
{"step": 47100, "lr": 1.0825767460079142e-05, "window_steps": 50, "elapsed_seconds": 3529.932405572, "steps_per_second": 1.7280784159977531, "eta_seconds": 1678.164586255541, "images_seen": 48230400, "images_per_second": 1769.5522979816992, "effective_batch_size": 1024, "loss": 0.37661463022232056, "edge_loss": 0.19403676688671112, "diag_loss": 0.18257789313793182, "jvp_rms": 2.4141016006469727, "grad_norm": 0.19179286062717438, "grad_norm_postclip": 0.19179286062717438, "clip_fraction": 0.0, "ema_updates": 47100, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.34063201977146995, "updates_nfe_5": 9, "loss_nfe_1": 0.40038707190089756, "updates_nfe_1": 9, "loss_nfe_2": 0.3866905212402344, "updates_nfe_2": 10, "loss_nfe_4": 0.37606876875673023, "updates_nfe_4": 14, "loss_nfe_3": 0.37871137633919716, "updates_nfe_3": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.505942459448033, "grad_norm": 0.05381970662946998, "grad_rms": 0.00015671772882359854, "grad_to_weight": 0.003260626090372893, "update_to_weight": 6.332156623985268e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.254501892925216, "grad_norm": 0.01609279312603895, "grad_rms": 4.686064170702597e-05, "grad_to_weight": 0.0009900514474112155, "update_to_weight": 9.635201814552247e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064937993505128, "grad_norm": 0.003603672397554724, "grad_rms": 6.451608512635706e-05, "grad_to_weight": 0.0007114938824869694, "update_to_weight": 2.8728806718788346e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.861841681328023, "grad_norm": 0.01415240449672563, "grad_rms": 4.528775235788213e-05, "grad_to_weight": 0.001100340437036413, "update_to_weight": 7.335215208487497e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045322805460887, "grad_norm": 0.021261349957631934, "grad_rms": 0.000483712460370036, "grad_to_weight": 0.004214071284917462, "update_to_weight": 2.2832290750499066e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642578927245903, "grad_norm": 0.0436807590099721, "grad_rms": 0.00010525514656285729, "grad_to_weight": 0.0026246388375819234, "update_to_weight": 6.078707989480695e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570260333807898, "grad_norm": 0.027003919424649477, "grad_rms": 3.725596542187263e-05, "grad_to_weight": 0.0010990489745643004, "update_to_weight": 8.504334386152223e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831956767378045, "grad_norm": 0.015572831956491208, "grad_rms": 1.8905055204829185e-05, "grad_to_weight": 0.0005050873700292725, "update_to_weight": 6.923973434002755e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.628525316926215, "grad_norm": 0.010381932162240894, "grad_rms": 9.319883732349545e-06, "grad_to_weight": 0.00029980867123920716, "update_to_weight": 8.384328512656966e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07749452951126, "grad_norm": 0.00897112381302404, "grad_rms": 7.216029433818171e-06, "grad_to_weight": 0.00024195604171375743, "update_to_weight": 8.738212158646797e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.66906699109104, "grad_norm": 0.0035101217554727383, "grad_rms": 6.3417066759604705e-06, "grad_to_weight": 0.00014228838718304102, "update_to_weight": 5.955747297054062e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.346075699023416, "grad_norm": 0.009018774930266019, "grad_rms": 7.2543581729749365e-06, "grad_to_weight": 0.00024149190407445824, "update_to_weight": 8.898636319747648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48501722540032, "grad_norm": 0.02265999147540694, "grad_rms": 1.0060251089638246e-05, "grad_to_weight": 0.00035693447786197796, "update_to_weight": 9.849561386076029e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.777600830422756, "grad_norm": 0.03214707107467462, "grad_rms": 2.3987212415888938e-05, "grad_to_weight": 0.0006330955096132488, "update_to_weight": 7.347558118138179e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.09969570586696, "grad_norm": 0.03293085840854033, "grad_rms": 2.8279721520701694e-05, "grad_to_weight": 0.0008212246459446882, "update_to_weight": 8.071698794563092e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20311958968092, "grad_norm": 0.05535965007352414, "grad_rms": 8.629784951839041e-05, "grad_to_weight": 0.0019628910162753544, "update_to_weight": 6.193028053294519e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41403212404568, "grad_norm": 0.006506653656840278, "grad_rms": 0.000553883141877058, "grad_to_weight": 0.0007733098187544967, "update_to_weight": 4.4926889558510835e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3192507746703053, "grad_norm": 0.09119242657980975, "grad_rms": 0.0021110715475922614, "grad_to_weight": 0.06912440631509176, "update_to_weight": 0.00012758151170817186, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47100}
{"step": 47150, "lr": 1.0797631506714433e-05, "window_steps": 50, "elapsed_seconds": 3557.425539719, "steps_per_second": 1.72877827837425, "eta_seconds": 1648.5630549917316, "images_seen": 48281600, "images_per_second": 1770.268957055232, "effective_batch_size": 1024, "loss": 0.38139593601226807, "edge_loss": 0.195024773478508, "diag_loss": 0.18637122213840485, "jvp_rms": 2.2822656631469727, "grad_norm": 0.2580885887145996, "grad_norm_postclip": 0.2580885887145996, "clip_fraction": 0.0, "ema_updates": 47150, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3774461881680922, "updates_nfe_4": 11, "loss_nfe_5": 0.4132108213379979, "updates_nfe_5": 16, "loss_nfe_1": 0.40112825731436413, "updates_nfe_1": 6, "loss_nfe_2": 0.36945311427116395, "updates_nfe_2": 5, "loss_nfe_3": 0.33770689740777016, "updates_nfe_3": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.505615203680023, "grad_norm": 0.008730494127274013, "grad_rms": 2.542234614086418e-05, "grad_to_weight": 0.000528940849495116, "update_to_weight": 5.840477177846055e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25884768899162, "grad_norm": 0.021481543592798816, "grad_rms": 6.255215671586632e-05, "grad_to_weight": 0.0013212217743661704, "update_to_weight": 5.677836547390754e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065140777167926, "grad_norm": 0.0031281247363067477, "grad_rms": 5.600241628799852e-05, "grad_to_weight": 0.0006175790316445611, "update_to_weight": 2.5352226151346673e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.863398092963106, "grad_norm": 0.012120005821095667, "grad_rms": 3.878406827114529e-05, "grad_to_weight": 0.0009422087175958498, "update_to_weight": 7.258155378864382e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045027167225365, "grad_norm": 0.019146761713312303, "grad_rms": 0.0004356039120244385, "grad_to_weight": 0.0037951751454774678, "update_to_weight": 1.6874645537663127e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642401056922438, "grad_norm": 0.03004106139965822, "grad_rms": 7.238830991473745e-05, "grad_to_weight": 0.0018050917831452322, "update_to_weight": 5.706368239940792e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57013861226227, "grad_norm": 0.021843075127368, "grad_rms": 3.0135805060570446e-05, "grad_to_weight": 0.0008890090313315015, "update_to_weight": 6.580534235168347e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832000318700434, "grad_norm": 0.013297083695375909, "grad_rms": 1.6142349832493522e-05, "grad_to_weight": 0.0004312754137885394, "update_to_weight": 5.3661619669688006e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.628856678264015, "grad_norm": 0.009617001609134848, "grad_rms": 8.633203863240152e-06, "grad_to_weight": 0.00027771640566958967, "update_to_weight": 6.315862353883362e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07827973524392, "grad_norm": 0.007728529352869905, "grad_rms": 6.2165338984032915e-06, "grad_to_weight": 0.0002084381855915426, "update_to_weight": 6.543589342435334e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669351517558685, "grad_norm": 0.003252973727053688, "grad_rms": 5.877119552738148e-06, "grad_to_weight": 0.00013186296059458018, "update_to_weight": 4.4648247501533235e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34714942686389, "grad_norm": 0.0077185492787338215, "grad_rms": 6.208506307856261e-06, "grad_to_weight": 0.00020667037236265883, "update_to_weight": 6.744368925084646e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48545370422464, "grad_norm": 0.019048910711716265, "grad_rms": 8.457056347613814e-06, "grad_to_weight": 0.00030005158032679624, "update_to_weight": 7.697372247472097e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.778002299993304, "grad_norm": 0.025939731597051895, "grad_rms": 1.9355475663218805e-05, "grad_to_weight": 0.000510845847061914, "update_to_weight": 6.142039374472815e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.100187336396104, "grad_norm": 0.029707573077535628, "grad_rms": 2.5511691291676877e-05, "grad_to_weight": 0.0007408337728779677, "update_to_weight": 7.231262754064242e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20316070807989, "grad_norm": 0.05736366763454006, "grad_rms": 8.942182890921162e-05, "grad_to_weight": 0.0020339446428820304, "update_to_weight": 5.645647195564493e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414108290410452, "grad_norm": 0.007577387458787741, "grad_rms": 0.0006450300560382503, "grad_to_weight": 0.0009005573968454482, "update_to_weight": 3.545911819650236e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3194401584801023, "grad_norm": 0.09970987616654646, "grad_rms": 0.0023082474113672513, "grad_to_weight": 0.07556983583204326, "update_to_weight": 9.663743529804865e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47150}
{"step": 47200, "lr": 1.0769978943369716e-05, "window_steps": 50, "elapsed_seconds": 3584.907925368, "steps_per_second": 1.7294725915069502, "eta_seconds": 1618.9906759726452, "images_seen": 48332800, "images_per_second": 1770.979933703117, "effective_batch_size": 1024, "loss": 0.36636537313461304, "edge_loss": 0.18756446242332458, "diag_loss": 0.17880094051361084, "jvp_rms": 2.194999933242798, "grad_norm": 0.21243688464164734, "grad_norm_postclip": 0.21243688464164734, "clip_fraction": 0.0, "ema_updates": 47200, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.34593938061824214, "updates_nfe_5": 13, "loss_nfe_3": 0.3943462123473485, "updates_nfe_3": 12, "loss_nfe_2": 0.3558608889579773, "updates_nfe_2": 11, "loss_nfe_1": 0.39845681687196094, "updates_nfe_1": 6, "loss_nfe_4": 0.34796149656176567, "updates_nfe_4": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.507429819866342, "grad_norm": 0.006219955134841147, "grad_rms": 1.8111901813735023e-05, "grad_to_weight": 0.00037679730901267033, "update_to_weight": 3.773447448412978e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.257106694193965, "grad_norm": 0.013098312867461407, "grad_rms": 3.814100768222051e-05, "grad_to_weight": 0.0008056976628036351, "update_to_weight": 4.394141748917319e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064846860598401, "grad_norm": 0.001799709020537506, "grad_rms": 3.221995996374618e-05, "grad_to_weight": 0.0003553333536178968, "update_to_weight": 2.319426685546581e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.86314601925576, "grad_norm": 0.007413189638711771, "grad_rms": 2.3722237208360726e-05, "grad_to_weight": 0.000576312328851312, "update_to_weight": 4.855036066223135e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045028490452995, "grad_norm": 0.012078185180840147, "grad_rms": 0.0002747882275712192, "grad_to_weight": 0.0023940767041645866, "update_to_weight": 2.0372011421198305e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642716454543354, "grad_norm": 0.06890779601090617, "grad_rms": 0.00016604336400829486, "grad_to_weight": 0.004140417593432879, "update_to_weight": 6.405794812306697e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570231766560525, "grad_norm": 0.013677527225748665, "grad_rms": 1.8870204482763825e-05, "grad_to_weight": 0.0005566706637404797, "update_to_weight": 7.61435528638991e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83198151245513, "grad_norm": 0.007088506908115694, "grad_rms": 8.605282250020143e-06, "grad_to_weight": 0.0002299076011463021, "update_to_weight": 7.125607972208551e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.629226812179546, "grad_norm": 0.003391946728545972, "grad_rms": 3.0449581679359033e-06, "grad_to_weight": 9.795040319389923e-05, "update_to_weight": 8.798541115983715e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07801470959155, "grad_norm": 0.002958797582310895, "grad_rms": 2.3799437938498826e-06, "grad_to_weight": 7.979924506436684e-05, "update_to_weight": 9.396923910732088e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669309457331035, "grad_norm": 0.0010454528823060022, "grad_rms": 1.8888106980292416e-06, "grad_to_weight": 4.237868449922592e-05, "update_to_weight": 6.296919025731273e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.347531843023106, "grad_norm": 0.0030064674713091146, "grad_rms": 2.4182876322905702e-06, "grad_to_weight": 8.049976324930165e-05, "update_to_weight": 9.202113833047367e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48590363960727, "grad_norm": 0.008084680969571297, "grad_rms": 3.589318231729237e-06, "grad_to_weight": 0.00012734608009150973, "update_to_weight": 9.433244316907646e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.778187407648275, "grad_norm": 0.018167542084562175, "grad_rms": 1.3556093183254502e-05, "grad_to_weight": 0.0003577824064240969, "update_to_weight": 6.832607499500681e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10037302809274, "grad_norm": 0.0286083007162352, "grad_rms": 2.4567679572719736e-05, "grad_to_weight": 0.0007134173214845995, "update_to_weight": 7.246031208492783e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20309037393996, "grad_norm": 0.06509374018466232, "grad_rms": 0.00010147191659601455, "grad_to_weight": 0.002308035726638306, "update_to_weight": 5.438076520928292e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414242033251755, "grad_norm": 0.002086244036582346, "grad_rms": 0.00017759288608972158, "grad_to_weight": 0.00024794200456058186, "update_to_weight": 3.0704532559456365e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3195375052756095, "grad_norm": 0.03707813372121597, "grad_rms": 0.0008583453261678028, "grad_to_weight": 0.028099340543921502, "update_to_weight": 7.945340600971227e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47200}
{"step": 47250, "lr": 1.0742810072449002e-05, "window_steps": 50, "elapsed_seconds": 3612.395191213, "steps_per_second": 1.7301540028629379, "eta_seconds": 1589.45388413372, "images_seen": 48384000, "images_per_second": 1771.6776989316484, "effective_batch_size": 1024, "loss": 0.36019420623779297, "edge_loss": 0.18642562627792358, "diag_loss": 0.17376859486103058, "jvp_rms": 2.2184765338897705, "grad_norm": 0.2089311182498932, "grad_norm_postclip": 0.2089311182498932, "clip_fraction": 0.0, "ema_updates": 47250, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4008333881696065, "updates_nfe_1": 9, "loss_nfe_3": 0.36318800426446474, "updates_nfe_3": 13, "loss_nfe_2": 0.37823321776730673, "updates_nfe_2": 14, "loss_nfe_5": 0.3135257959365845, "updates_nfe_5": 10, "loss_nfe_4": 0.3125605806708336, "updates_nfe_4": 4, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.506351057824396, "grad_norm": 0.020712613234760077, "grad_rms": 6.03131066191242e-05, "grad_to_weight": 0.001254826894339062, "update_to_weight": 2.9186422285733936e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.256739700609078, "grad_norm": 0.02084615936256283, "grad_rms": 6.070197989906462e-05, "grad_to_weight": 0.0012823087375743493, "update_to_weight": 4.572507543291683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06480637747954, "grad_norm": 0.003098307181903023, "grad_rms": 5.546859643259884e-05, "grad_to_weight": 0.0006117326016014201, "update_to_weight": 2.4842899012511615e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.863104500759563, "grad_norm": 0.014234995588790092, "grad_rms": 4.555204419078217e-05, "grad_to_weight": 0.0011066531868686536, "update_to_weight": 4.4654777772146116e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.045005239402689, "grad_norm": 0.06750541397078906, "grad_rms": 0.0015358013458777128, "grad_to_weight": 0.013380642985968725, "update_to_weight": 1.6183739531250722e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642694450252133, "grad_norm": 0.15629653237780164, "grad_rms": 0.0003766192437026151, "grad_to_weight": 0.009391299758882119, "update_to_weight": 5.67090475699277e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570446641128264, "grad_norm": 0.04608862108855363, "grad_rms": 6.358617971765896e-05, "grad_to_weight": 0.001875774655695769, "update_to_weight": 7.481550661779028e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83173109135852, "grad_norm": 0.03232790984950555, "grad_rms": 3.924532943456659e-05, "grad_to_weight": 0.001048527238179188, "update_to_weight": 6.882856688995685e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.629008257393465, "grad_norm": 0.02249149260394875, "grad_rms": 2.0190663236866782e-05, "grad_to_weight": 0.0006494986063930001, "update_to_weight": 8.329530511766156e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07809207752437, "grad_norm": 0.01788704754765051, "grad_rms": 1.4387658032382274e-05, "grad_to_weight": 0.000482415532877246, "update_to_weight": 8.892951952856656e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.669608825625406, "grad_norm": 0.007105352654407804, "grad_rms": 1.2837179306745383e-05, "grad_to_weight": 0.0002880204832038991, "update_to_weight": 6.370572460547174e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34769853602093, "grad_norm": 0.015837951418063385, "grad_rms": 1.2739443350253245e-05, "grad_to_weight": 0.00042406766785878583, "update_to_weight": 8.497580475679512e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.486872720370314, "grad_norm": 0.02670751585924533, "grad_rms": 1.1857211677070076e-05, "grad_to_weight": 0.00042067776714848314, "update_to_weight": 9.108134195075371e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.778391746483805, "grad_norm": 0.037859088269626646, "grad_rms": 2.8249354041801e-05, "grad_to_weight": 0.0007455747802853215, "update_to_weight": 6.603846805289656e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10063025533389, "grad_norm": 0.04207472811905801, "grad_rms": 3.61321159474427e-05, "grad_to_weight": 0.0010492285994298442, "update_to_weight": 7.058368925915713e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203140148887897, "grad_norm": 0.10265026420549041, "grad_rms": 0.00016001721545680472, "grad_to_weight": 0.003639675002981471, "update_to_weight": 5.127003126526168e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414416575796277, "grad_norm": 0.010650017865303733, "grad_rms": 0.0009065897260537149, "grad_to_weight": 0.001265687022905197, "update_to_weight": 4.034726773693373e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3197296029632126, "grad_norm": 0.1127131828076001, "grad_rms": 0.002609269236359718, "grad_to_weight": 0.08540627000752515, "update_to_weight": 0.00015082539994170431, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47250}
{"step": 47300, "lr": 1.07161251910667e-05, "window_steps": 50, "elapsed_seconds": 3639.871154796, "steps_per_second": 1.7308304970352966, "eta_seconds": 1559.9447806268572, "images_seen": 48435200, "images_per_second": 1772.3704289641437, "effective_batch_size": 1024, "loss": 0.353771835565567, "edge_loss": 0.18247324228286743, "diag_loss": 0.17129866778850555, "jvp_rms": 1.8741406202316284, "grad_norm": 0.21455861628055573, "grad_norm_postclip": 0.21455861628055573, "clip_fraction": 0.0, "ema_updates": 47300, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3635379791259766, "updates_nfe_4": 10, "loss_nfe_3": 0.3467588394880295, "updates_nfe_3": 10, "loss_nfe_5": 0.3295239818592866, "updates_nfe_5": 12, "loss_nfe_2": 0.34297777546776664, "updates_nfe_2": 9, "loss_nfe_1": 0.39383767379654777, "updates_nfe_1": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.506775360817752, "grad_norm": 0.029608059302818325, "grad_rms": 8.621577669974383e-05, "grad_to_weight": 0.0017936912967931448, "update_to_weight": 4.4004964770187785e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25656324038608, "grad_norm": 0.010848318285180935, "grad_rms": 3.15892432285786e-05, "grad_to_weight": 0.0006673192928152566, "update_to_weight": 4.8351492899049855e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064823888821829, "grad_norm": 0.001314108187183918, "grad_rms": 2.3526310473484973e-05, "grad_to_weight": 0.00025945782440415783, "update_to_weight": 2.2681692856192323e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.863530355544228, "grad_norm": 0.010912040045883099, "grad_rms": 3.491857284262776e-05, "grad_to_weight": 0.0008482927893259078, "update_to_weight": 5.739102152829143e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.044723382782158, "grad_norm": 0.018694562154773286, "grad_rms": 0.0004253160158431033, "grad_to_weight": 0.003705765556656401, "update_to_weight": 2.1796319472056796e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642697200790128, "grad_norm": 0.05520847710592215, "grad_rms": 0.00013303286117279626, "grad_to_weight": 0.003317279431323252, "update_to_weight": 6.856913374869623e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.570554698726788, "grad_norm": 0.02440525611359978, "grad_rms": 3.367071881610835e-05, "grad_to_weight": 0.0009932724927396298, "update_to_weight": 8.517179217021907e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831683580407084, "grad_norm": 0.01227859085967152, "grad_rms": 1.490592325712753e-05, "grad_to_weight": 0.0003982458767666621, "update_to_weight": 7.66387150103702e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62913868541564, "grad_norm": 0.008590027849047997, "grad_rms": 7.71128722088392e-06, "grad_to_weight": 0.00024805779684799843, "update_to_weight": 9.035493935226178e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07820401382226, "grad_norm": 0.007114589116201472, "grad_rms": 5.722703815267772e-06, "grad_to_weight": 0.0001918806292114162, "update_to_weight": 9.17980778421051e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670301564773688, "grad_norm": 0.0026601704438487914, "grad_rms": 4.806106978097282e-06, "grad_to_weight": 0.00010782885798393338, "update_to_weight": 5.94573225877544e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34845027951294, "grad_norm": 0.007277022229204705, "grad_rms": 5.853358808876977e-06, "grad_to_weight": 0.00019484134347594152, "update_to_weight": 9.460938175725771e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48773988129027, "grad_norm": 0.017599233066871625, "grad_rms": 7.81344970186548e-06, "grad_to_weight": 0.00027720679771840625, "update_to_weight": 9.835991876787771e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77840617037054, "grad_norm": 0.025614637465362907, "grad_rms": 1.9112899847404404e-05, "grad_to_weight": 0.0005044395718018653, "update_to_weight": 7.050013377719838e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10088291480469, "grad_norm": 0.025813368366329168, "grad_rms": 2.2167501978076635e-05, "grad_to_weight": 0.0006437107238055158, "update_to_weight": 7.762454980364512e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20307089676251, "grad_norm": 0.06270755636382326, "grad_rms": 9.775219416242802e-05, "grad_to_weight": 0.0022234300865095363, "update_to_weight": 6.320477358137613e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414455564025806, "grad_norm": 0.00863529338779111, "grad_rms": 0.0007350849891375023, "grad_to_weight": 0.0010262450519923653, "update_to_weight": 2.7887366100828196e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3198218703734195, "grad_norm": 0.10978730859769102, "grad_rms": 0.0025415363112909098, "grad_to_weight": 0.08318342881121428, "update_to_weight": 8.873082597839187e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47300}
{"step": 47350, "lr": 1.068992459104441e-05, "window_steps": 50, "elapsed_seconds": 3667.373380797, "steps_per_second": 1.7314844551279387, "eta_seconds": 1530.4786549782757, "images_seen": 48486400, "images_per_second": 1773.0400820510092, "effective_batch_size": 1024, "loss": 0.352520614862442, "edge_loss": 0.1848151683807373, "diag_loss": 0.1677054762840271, "jvp_rms": 1.7298437356948853, "grad_norm": 0.2816544473171234, "grad_norm_postclip": 0.2816544473171234, "clip_fraction": 0.0, "ema_updates": 47350, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4095532321012937, "updates_nfe_1": 13, "loss_nfe_4": 0.32803068657716117, "updates_nfe_4": 15, "loss_nfe_3": 0.32207908374922617, "updates_nfe_3": 7, "loss_nfe_2": 0.3446666896343231, "updates_nfe_2": 7, "loss_nfe_5": 0.33926994539797306, "updates_nfe_5": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.504684245715858, "grad_norm": 0.31396179435044086, "grad_rms": 0.000914226078687701, "grad_to_weight": 0.019022587144127665, "update_to_weight": 4.7155419350069994e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25790917048386, "grad_norm": 0.22444006301948258, "grad_rms": 0.0006535475411562092, "grad_to_weight": 0.013804977052458395, "update_to_weight": 4.4555884628847093e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.064869832270784, "grad_norm": 0.008768984258771585, "grad_rms": 0.0001569900014480996, "grad_to_weight": 0.0017313345750566108, "update_to_weight": 2.5439325584265547e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.863781235538067, "grad_norm": 0.102770349237953, "grad_rms": 0.0003288655385096083, "grad_to_weight": 0.00798912445386081, "update_to_weight": 6.177215985210015e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.044340176188592, "grad_norm": 0.05244155667104051, "grad_rms": 0.0011930867256092607, "grad_to_weight": 0.010396118191748193, "update_to_weight": 3.37416640550833e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.642855814378827, "grad_norm": 0.2645132275439532, "grad_rms": 0.0006373831216302163, "grad_to_weight": 0.01589349991937221, "update_to_weight": 6.828179348894858e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.571099947636018, "grad_norm": 0.029662200930230486, "grad_rms": 4.092346428735686e-05, "grad_to_weight": 0.0012071987413442712, "update_to_weight": 6.773022051827514e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831705356268913, "grad_norm": 0.011681006665671413, "grad_rms": 1.4180469967149896e-05, "grad_to_weight": 0.00037886346313621445, "update_to_weight": 4.916811865544863e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62931846377611, "grad_norm": 0.0034971463442222115, "grad_rms": 3.1393960983200277e-06, "grad_to_weight": 0.00010098802111512506, "update_to_weight": 5.06797527761998e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07900731145759, "grad_norm": 0.002656364071144028, "grad_rms": 2.136677825857692e-06, "grad_to_weight": 7.16406469253884e-05, "update_to_weight": 5.574803712054305e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670114774525903, "grad_norm": 0.001038116930390759, "grad_rms": 1.8755568970284962e-06, "grad_to_weight": 4.2079939225200016e-05, "update_to_weight": 3.487297460861286e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34883104816083, "grad_norm": 0.002804741896690648, "grad_rms": 2.2560272829364065e-06, "grad_to_weight": 7.509584150234767e-05, "update_to_weight": 5.973410845549091e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48837053633819, "grad_norm": 0.008730611218606382, "grad_rms": 3.876088882050888e-06, "grad_to_weight": 0.000137515125130032, "update_to_weight": 7.21881612669666e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77888215633419, "grad_norm": 0.0311778070979736, "grad_rms": 2.326397573774214e-05, "grad_to_weight": 0.0006139915999329351, "update_to_weight": 6.893321064299817e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.101094477835794, "grad_norm": 0.09571174660598271, "grad_rms": 8.21934705344698e-05, "grad_to_weight": 0.002386761455074085, "update_to_weight": 9.258709488474886e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.2032710782228, "grad_norm": 0.2806091557513626, "grad_rms": 0.00043742990904660717, "grad_to_weight": 0.009949525180007768, "update_to_weight": 8.343705220086782e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414182189260137, "grad_norm": 0.004627705185392445, "grad_rms": 0.0003939364261491444, "grad_to_weight": 0.0005499887073159942, "update_to_weight": 7.607380680911701e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3195639299645041, "grad_norm": 0.0752584638089224, "grad_rms": 0.001742206097821873, "grad_to_weight": 0.057032828876238555, "update_to_weight": 0.0002670834532369075, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47350}
{"step": 47400, "lr": 1.066420855890769e-05, "window_steps": 50, "elapsed_seconds": 3694.900282697, "steps_per_second": 1.73211711016149, "eta_seconds": 1501.0532398456562, "images_seen": 48537600, "images_per_second": 1773.6879208053658, "effective_batch_size": 1024, "loss": 0.39447295665740967, "edge_loss": 0.20622695982456207, "diag_loss": 0.18824604153633118, "jvp_rms": 2.711249828338623, "grad_norm": 0.34248384833335876, "grad_norm_postclip": 0.34248384833335876, "clip_fraction": 0.0, "ema_updates": 47400, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4052087515592575, "updates_nfe_1": 14, "loss_nfe_4": 0.44146180694753473, "updates_nfe_4": 11, "loss_nfe_5": 0.30876172582308453, "updates_nfe_5": 6, "loss_nfe_2": 0.35145580768585205, "updates_nfe_2": 6, "loss_nfe_3": 0.4025647227580731, "updates_nfe_3": 13, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50708225727483, "grad_norm": 0.07533758810553867, "grad_rms": 0.00021937569790621592, "grad_to_weight": 0.004563955454473892, "update_to_weight": 4.1905640847262e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.256336091960133, "grad_norm": 0.03588626490196249, "grad_rms": 0.00010449729817587016, "grad_to_weight": 0.002207524789039684, "update_to_weight": 4.5159307175320417e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065221926015685, "grad_norm": 0.004478844693884435, "grad_rms": 8.018418259508086e-05, "grad_to_weight": 0.0008842346415031621, "update_to_weight": 2.9585361093049513e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.86399178116569, "grad_norm": 0.029664366305940236, "grad_rms": 9.492609368448643e-05, "grad_to_weight": 0.002306000097836828, "update_to_weight": 6.302266294149667e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.044929814526273, "grad_norm": 0.032823985507875915, "grad_rms": 0.0007467715277159884, "grad_to_weight": 0.006506331448529406, "update_to_weight": 2.666601323774604e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.643488420013078, "grad_norm": 0.10907514801438305, "grad_rms": 0.0002628324450131063, "grad_to_weight": 0.006553622970243719, "update_to_weight": 7.662648549598163e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.571537131643165, "grad_norm": 0.041794749399753875, "grad_rms": 5.7662138372095735e-05, "grad_to_weight": 0.0017009415884662218, "update_to_weight": 9.319900887687757e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83149452568753, "grad_norm": 0.030476327673663272, "grad_rms": 3.699755180822514e-05, "grad_to_weight": 0.0009884803880743328, "update_to_weight": 8.110225061529567e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62852884206352, "grad_norm": 0.021859092095651126, "grad_rms": 1.9622955885527236e-05, "grad_to_weight": 0.0006312451850134255, "update_to_weight": 9.604700914580003e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07844434649666, "grad_norm": 0.01865465751497444, "grad_rms": 1.5005094178995273e-05, "grad_to_weight": 0.0005031132735949591, "update_to_weight": 0.00010202456824646948, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670490827622, "grad_norm": 0.0065893871375337365, "grad_rms": 1.1904988861264672e-05, "grad_to_weight": 0.0002670959075591643, "update_to_weight": 6.774018575596235e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.34906800523107, "grad_norm": 0.017454269791428768, "grad_rms": 1.4039548130848635e-05, "grad_to_weight": 0.00046732812152057294, "update_to_weight": 0.00010372975992827436, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.4880552095973, "grad_norm": 0.03421375586316156, "grad_rms": 1.5189722161923572e-05, "grad_to_weight": 0.0005389006758863449, "update_to_weight": 0.00011303828536467377, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77890138799652, "grad_norm": 0.04426235434500255, "grad_rms": 3.302728554133659e-05, "grad_to_weight": 0.0008716682152455075, "update_to_weight": 8.421687639292075e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.101817436604414, "grad_norm": 0.04515241985345983, "grad_rms": 3.8775116141840205e-05, "grad_to_weight": 0.0011259444768267608, "update_to_weight": 9.134984838993445e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203542673500035, "grad_norm": 0.10346748024182399, "grad_rms": 0.00016129113945079435, "grad_to_weight": 0.003668598709021109, "update_to_weight": 7.023291127196654e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414470524577549, "grad_norm": 0.006016240800574692, "grad_rms": 0.000512136427210645, "grad_to_weight": 0.0007149874472793093, "update_to_weight": 5.09881308309555e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3198929972338795, "grad_norm": 0.07324003216433184, "grad_rms": 0.0016954801384909653, "grad_to_weight": 0.05548937097008782, "update_to_weight": 0.00011568308182834013, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47400}
{"step": 47450, "lr": 1.0638977375882945e-05, "window_steps": 50, "elapsed_seconds": 3722.4434354610003, "steps_per_second": 1.7327328438507783, "eta_seconds": 1471.663683786907, "images_seen": 48588800, "images_per_second": 1774.318432103197, "effective_batch_size": 1024, "loss": 0.3724244236946106, "edge_loss": 0.191632479429245, "diag_loss": 0.1807919442653656, "jvp_rms": 2.060429573059082, "grad_norm": 0.2099541276693344, "grad_norm_postclip": 0.2099541276693344, "clip_fraction": 0.0, "ema_updates": 47450, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.365720234811306, "updates_nfe_4": 8, "loss_nfe_5": 0.3794497239092986, "updates_nfe_5": 12, "loss_nfe_3": 0.37665153633464465, "updates_nfe_3": 11, "loss_nfe_2": 0.34692856669425964, "updates_nfe_2": 11, "loss_nfe_1": 0.39783527702093124, "updates_nfe_1": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50693805362612, "grad_norm": 0.49707112190465474, "grad_rms": 0.0014474225551805602, "grad_to_weight": 0.030112860440247545, "update_to_weight": 4.036469383028002e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25697529079345, "grad_norm": 0.3563577831917019, "grad_rms": 0.0010376790571324944, "grad_to_weight": 0.021920300475175864, "update_to_weight": 4.7799949293744186e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065331503139191, "grad_norm": 0.013141678033535232, "grad_rms": 0.00023527377774129664, "grad_to_weight": 0.002594435926926165, "update_to_weight": 2.3724813087554752e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.863812076239784, "grad_norm": 0.16490264197655458, "grad_rms": 0.000527689129767493, "grad_to_weight": 0.012819111551010563, "update_to_weight": 5.4176397972549525e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0448035367702255, "grad_norm": 0.05211122878407493, "grad_rms": 0.0011855715059618085, "grad_to_weight": 0.010329684477155572, "update_to_weight": 3.137694506339999e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.643261051480696, "grad_norm": 0.2948758676782175, "grad_rms": 0.0007105463223117267, "grad_to_weight": 0.017717433306256008, "update_to_weight": 6.263444324913738e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.5713545585861, "grad_norm": 0.03266882328403723, "grad_rms": 4.50715517071336e-05, "grad_to_weight": 0.0013295491384548675, "update_to_weight": 5.792489695621092e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831594497175804, "grad_norm": 0.013806242124601173, "grad_rms": 1.676045630403324e-05, "grad_to_weight": 0.0004477952681255533, "update_to_weight": 4.49123946529391e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62917746121939, "grad_norm": 0.0036721713445942916, "grad_rms": 3.296516432784834e-06, "grad_to_weight": 0.0001060426961832025, "update_to_weight": 4.5666734290858375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.079125829307586, "grad_norm": 0.0029239304609500307, "grad_rms": 2.351898012824407e-06, "grad_to_weight": 7.885651011327071e-05, "update_to_weight": 5.046360706058167e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670358467482036, "grad_norm": 0.0013211352398119303, "grad_rms": 2.386883633623012e-06, "grad_to_weight": 5.3551521821351594e-05, "update_to_weight": 3.179086812293429e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.350283816483575, "grad_norm": 0.002951807345979181, "grad_rms": 2.3743211146660593e-06, "grad_to_weight": 7.903038596661152e-05, "update_to_weight": 5.5539075932538494e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.489229986027944, "grad_norm": 0.009384358411621685, "grad_rms": 4.166329984657572e-06, "grad_to_weight": 0.00014781024141711748, "update_to_weight": 6.68640336034355e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.779103320011224, "grad_norm": 0.04493435655182019, "grad_rms": 3.3528714105121974e-05, "grad_to_weight": 0.000884898582565405, "update_to_weight": 5.8033321444878546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.1019468067987, "grad_norm": 0.1595596564387308, "grad_rms": 0.00013702353561654857, "grad_to_weight": 0.003978850633049061, "update_to_weight": 7.63384744729307e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203183431379912, "grad_norm": 0.4089681037354976, "grad_rms": 0.0006375233193691451, "grad_to_weight": 0.014500778067502282, "update_to_weight": 6.0061884188317486e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414613328505459, "grad_norm": 0.0048673190236690434, "grad_rms": 0.00041433371061845695, "grad_to_weight": 0.0005784364454609515, "update_to_weight": 3.4748931365669226e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3201482095020902, "grad_norm": 0.08026185706296524, "grad_rms": 0.0018580328340561676, "grad_to_weight": 0.06079761081768006, "update_to_weight": 7.253765107167811e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47450}
{"step": 47500, "lr": 1.0614231317894356e-05, "window_steps": 50, "elapsed_seconds": 3749.979885635, "steps_per_second": 1.7333426306896917, "eta_seconds": 1442.2999560134615, "images_seen": 48640000, "images_per_second": 1774.9428538262443, "effective_batch_size": 1024, "loss": 0.3614285886287689, "edge_loss": 0.18740272521972656, "diag_loss": 0.17402590811252594, "jvp_rms": 1.9802733659744263, "grad_norm": 0.2611817419528961, "grad_norm_postclip": 0.24795085191726685, "clip_fraction": 0.019999999552965164, "ema_updates": 47500, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.29390957123703426, "updates_nfe_5": 9, "loss_nfe_1": 0.4014132857322693, "updates_nfe_1": 10, "loss_nfe_3": 0.36715860664844513, "updates_nfe_3": 9, "loss_nfe_4": 0.374360849459966, "updates_nfe_4": 15, "loss_nfe_2": 0.35603874921798706, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50702124819258, "grad_norm": 0.07311486303337834, "grad_rms": 0.0002129033396024734, "grad_to_weight": 0.004429319011228872, "update_to_weight": 5.570487646195619e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.258035873645685, "grad_norm": 0.02337297607762435, "grad_rms": 6.805982336454929e-05, "grad_to_weight": 0.0014376260613074424, "update_to_weight": 5.780387357075504e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065166006879349, "grad_norm": 0.003478012837867137, "grad_rms": 6.226641813241888e-05, "grad_to_weight": 0.0006866532771370987, "update_to_weight": 2.6607651434228197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.865322587038177, "grad_norm": 0.017590674748414324, "grad_rms": 5.629023124646794e-05, "grad_to_weight": 0.0013672937176201817, "update_to_weight": 7.169136419040567e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.044137690698442, "grad_norm": 0.035449045551286786, "grad_rms": 0.0008064937116200048, "grad_to_weight": 0.007027771192022773, "update_to_weight": 2.489936240044124e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.643666278715926, "grad_norm": 0.05570096332079363, "grad_rms": 0.00013421957838882838, "grad_to_weight": 0.0033466762904291424, "update_to_weight": 5.439290465050162e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57208608463639, "grad_norm": 0.02561922737773455, "grad_rms": 3.5345574629759894e-05, "grad_to_weight": 0.0010426150750689777, "update_to_weight": 6.0475509767828006e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83180928630669, "grad_norm": 0.014197721577024485, "grad_rms": 1.7235703239227586e-05, "grad_to_weight": 0.0004604894070660365, "update_to_weight": 4.970870209626533e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62859229447358, "grad_norm": 0.0101206636014903, "grad_rms": 9.085342360756895e-06, "grad_to_weight": 0.0002922632117247651, "update_to_weight": 5.815381231833249e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.07920154884676, "grad_norm": 0.00824816282750551, "grad_rms": 6.634507223252965e-06, "grad_to_weight": 0.00022244715320095788, "update_to_weight": 6.203461649817511e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67039805189739, "grad_norm": 0.003509605966463366, "grad_rms": 6.340774804409569e-06, "grad_to_weight": 0.00014225980298657743, "update_to_weight": 3.972346261703156e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35026420693868, "grad_norm": 0.00835095913682324, "grad_rms": 6.717192649202356e-06, "grad_to_weight": 0.00022358500840997678, "update_to_weight": 6.446630201140731e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.48918384137372, "grad_norm": 0.01988358320825008, "grad_rms": 8.827622016266372e-06, "grad_to_weight": 0.0003131806396807488, "update_to_weight": 7.714534626520155e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77913697526894, "grad_norm": 0.03094356015123452, "grad_rms": 2.3089187457461445e-05, "grad_to_weight": 0.0006093754639096174, "update_to_weight": 6.224670730309263e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10228316735065, "grad_norm": 0.036925203922753995, "grad_rms": 3.170990780367243e-05, "grad_to_weight": 0.0009207756019441986, "update_to_weight": 7.699838393793812e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203544837595327, "grad_norm": 0.08160931920136033, "grad_rms": 0.00012721736388116137, "grad_to_weight": 0.0028935837559176288, "update_to_weight": 7.136438285225301e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414232966007656, "grad_norm": 0.009192731839427757, "grad_rms": 0.0007825373013826762, "grad_to_weight": 0.0010925216685305874, "update_to_weight": 7.451938196479899e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.319718131184797, "grad_norm": 0.1083010545388326, "grad_rms": 0.0025071300697440463, "grad_to_weight": 0.08206377708973633, "update_to_weight": 0.00018096569055851833, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47500}
{"step": 47550, "lr": 1.0589970655560846e-05, "window_steps": 50, "elapsed_seconds": 3779.677264295, "steps_per_second": 1.7329521919437563, "eta_seconds": 1413.7724118355343, "images_seen": 48691200, "images_per_second": 1774.5430445504064, "effective_batch_size": 1024, "loss": 0.3565216064453125, "edge_loss": 0.18587639927864075, "diag_loss": 0.1706451028585434, "jvp_rms": 1.9127733707427979, "grad_norm": 0.2159881740808487, "grad_norm_postclip": 0.2159881740808487, "clip_fraction": 0.0, "ema_updates": 47550, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.29935001499123043, "updates_nfe_5": 9, "loss_nfe_1": 0.40102026554254383, "updates_nfe_1": 13, "loss_nfe_2": 0.35734513672915375, "updates_nfe_2": 11, "loss_nfe_3": 0.3596044567724069, "updates_nfe_3": 12, "loss_nfe_4": 0.3345226407051086, "updates_nfe_4": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.507934512090504, "grad_norm": 0.04908989831498411, "grad_rms": 0.00014294498899949636, "grad_to_weight": 0.002973715353609526, "update_to_weight": 5.57111670954607e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25647031641363, "grad_norm": 0.015899809216095188, "grad_rms": 4.6298691411122206e-05, "grad_to_weight": 0.0009780603603748881, "update_to_weight": 4.265548330444959e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0659375234938855, "grad_norm": 0.0031269644382650105, "grad_rms": 5.5981643620848354e-05, "grad_to_weight": 0.0006172528626267778, "update_to_weight": 2.4848159607214208e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.864463272255206, "grad_norm": 0.015604538017723458, "grad_rms": 4.99345855730255e-05, "grad_to_weight": 0.0012129956522459644, "update_to_weight": 6.540971754855386e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.04462772564843, "grad_norm": 0.0328167176090541, "grad_rms": 0.0007466061772924346, "grad_to_weight": 0.006505280348479209, "update_to_weight": 2.4434940737038514e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.643229879746244, "grad_norm": 0.07438245712890931, "grad_rms": 0.00017923535680827905, "grad_to_weight": 0.004469232094151872, "update_to_weight": 7.125783001568942e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.571683685904723, "grad_norm": 0.036250981207631726, "grad_rms": 5.0013676946009214e-05, "grad_to_weight": 0.0014753153129847061, "update_to_weight": 8.407987244976995e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83203694135574, "grad_norm": 0.020184581588868575, "grad_rms": 2.4503611821540343e-05, "grad_to_weight": 0.0006546626039421521, "update_to_weight": 7.031402801447726e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63030370311979, "grad_norm": 0.014568586564462896, "grad_rms": 1.3078252757158925e-05, "grad_to_weight": 0.00042068896332406213, "update_to_weight": 8.062581193734266e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08045912901767, "grad_norm": 0.011155413468829567, "grad_rms": 8.97298862608732e-06, "grad_to_weight": 0.00030084345584868423, "update_to_weight": 8.283695941764432e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670769152702505, "grad_norm": 0.004157035387330411, "grad_rms": 7.510479950427444e-06, "grad_to_weight": 0.00016850043716107804, "update_to_weight": 5.533945893005629e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.351499588160785, "grad_norm": 0.01095961296221599, "grad_rms": 8.815494175188217e-06, "grad_to_weight": 0.00029341828529127735, "update_to_weight": 7.941813526642739e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.491223787598194, "grad_norm": 0.024689233489304813, "grad_rms": 1.0961164234447325e-05, "grad_to_weight": 0.0003888605702718773, "update_to_weight": 8.79938324071412e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77948554626908, "grad_norm": 0.03542453172921641, "grad_rms": 2.6432758534930143e-05, "grad_to_weight": 0.000697615018114715, "update_to_weight": 7.389501688070526e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.102584519657206, "grad_norm": 0.03160324530429606, "grad_rms": 2.7139619783617144e-05, "grad_to_weight": 0.0007880600635304441, "update_to_weight": 8.497676920084377e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203133656508367, "grad_norm": 0.05860359983306261, "grad_rms": 9.13547074974793e-05, "grad_to_weight": 0.0020779109352459775, "update_to_weight": 7.031861289951348e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414972820997404, "grad_norm": 0.005036619970709379, "grad_rms": 0.00042874556430162964, "grad_to_weight": 0.0005985307472582427, "update_to_weight": 4.732545518697201e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3205502087293102, "grad_norm": 0.07734551487941639, "grad_rms": 0.0017905205719348728, "grad_to_weight": 0.05857067332096486, "update_to_weight": 0.00014956625986820655, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47550}
{"step": 47600, "lr": 1.0566195654193134e-05, "window_steps": 50, "elapsed_seconds": 3807.1574823700003, "steps_per_second": 1.733576830105652, "eta_seconds": 1384.4209026800002, "images_seen": 48742400, "images_per_second": 1775.1826740281876, "effective_batch_size": 1024, "loss": 0.3484290540218353, "edge_loss": 0.18260090053081512, "diag_loss": 0.16582822799682617, "jvp_rms": 1.810585856437683, "grad_norm": 0.21249160170555115, "grad_norm_postclip": 0.21249160170555115, "clip_fraction": 0.0, "ema_updates": 47600, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3466221872303221, "updates_nfe_3": 9, "loss_nfe_4": 0.3145378033320109, "updates_nfe_4": 12, "loss_nfe_1": 0.39955063802855356, "updates_nfe_1": 14, "loss_nfe_2": 0.36108240857720375, "updates_nfe_2": 8, "loss_nfe_5": 0.2921476342848369, "updates_nfe_5": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.506829900062865, "grad_norm": 0.01891553248948926, "grad_rms": 5.508018302014552e-05, "grad_to_weight": 0.0011459215733129486, "update_to_weight": 3.7803379719250655e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.257239973700948, "grad_norm": 0.006742996454831819, "grad_rms": 1.963494704908291e-05, "grad_to_weight": 0.0004147688331930787, "update_to_weight": 4.243557234851702e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065491345929198, "grad_norm": 0.0022130509276219903, "grad_rms": 3.961996715692179e-05, "grad_to_weight": 0.00043688771265999174, "update_to_weight": 2.1552594704460458e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.86488137253866, "grad_norm": 0.010020970841585857, "grad_rms": 3.2067147739050116e-05, "grad_to_weight": 0.0007789400112912502, "update_to_weight": 5.33869619674832e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043961856368076, "grad_norm": 0.04403032451229613, "grad_rms": 0.0010017245679684028, "grad_to_weight": 0.00872931353687919, "update_to_weight": 3.075561722475494e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64336740166475, "grad_norm": 0.12467324868002522, "grad_rms": 0.0003004183388683289, "grad_to_weight": 0.0074908668222726845, "update_to_weight": 7.235339685329322e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.572185441333485, "grad_norm": 0.03883914421049698, "grad_rms": 5.358443679848087e-05, "grad_to_weight": 0.0015806141583631665, "update_to_weight": 7.672487045014388e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83183106207974, "grad_norm": 0.023531267974711768, "grad_rms": 2.8566411128332022e-05, "grad_to_weight": 0.0007632134441620309, "update_to_weight": 6.369485616019812e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.62942597783982, "grad_norm": 0.016763476953025457, "grad_rms": 1.5048610770193678e-05, "grad_to_weight": 0.0004840818604314376, "update_to_weight": 7.105055152598537e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08017930465851, "grad_norm": 0.013452569254854552, "grad_rms": 1.0820733023724203e-05, "grad_to_weight": 0.0003627967692476735, "update_to_weight": 7.59129247233648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670678852020494, "grad_norm": 0.004975396141635884, "grad_rms": 8.989005261075568e-06, "grad_to_weight": 0.00020167244571903648, "update_to_weight": 5.181176440751272e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.351473442949676, "grad_norm": 0.012428889602307034, "grad_rms": 9.99732420031018e-06, "grad_to_weight": 0.00033275500152064457, "update_to_weight": 7.97097966251315e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49109497106661, "grad_norm": 0.023874246798862443, "grad_rms": 1.0599338381623773e-05, "grad_to_weight": 0.00037602512304665917, "update_to_weight": 8.703456479367466e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7796874759608, "grad_norm": 0.034079706897011244, "grad_rms": 2.5429289234808403e-05, "grad_to_weight": 0.0006711287247119165, "update_to_weight": 7.261380329409142e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10264539863611, "grad_norm": 0.03371806921008712, "grad_rms": 2.895574708826048e-05, "grad_to_weight": 0.0008407941390129307, "update_to_weight": 8.531942849437037e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203510212050734, "grad_norm": 0.06990889103709032, "grad_rms": 0.00010897805442599175, "grad_to_weight": 0.002478730147824642, "update_to_weight": 6.73079292123534e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414472791325492, "grad_norm": 0.0029154483649187833, "grad_rms": 0.000248179446072703, "grad_to_weight": 0.0003464802177415475, "update_to_weight": 3.34287363423655e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.320075245143117, "grad_norm": 0.05194959780902339, "grad_rms": 0.0012026143174017879, "grad_to_weight": 0.03935351261237478, "update_to_weight": 0.00011684951128579337, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47600}
{"step": 47650, "lr": 1.0542906573790817e-05, "window_steps": 50, "elapsed_seconds": 3834.624770238, "steps_per_second": 1.7341983631914162, "eta_seconds": 1355.0929639186918, "images_seen": 48793600, "images_per_second": 1775.8191239080102, "effective_batch_size": 1024, "loss": 0.36339160799980164, "edge_loss": 0.18856805562973022, "diag_loss": 0.174823597073555, "jvp_rms": 2.088515520095825, "grad_norm": 0.2673896551132202, "grad_norm_postclip": 0.2673896551132202, "clip_fraction": 0.0, "ema_updates": 47650, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4041801772334359, "updates_nfe_1": 11, "loss_nfe_5": 0.36306485669179395, "updates_nfe_5": 11, "loss_nfe_2": 0.35697106804166523, "updates_nfe_2": 7, "loss_nfe_4": 0.3239241898059845, "updates_nfe_4": 10, "loss_nfe_3": 0.36289527199485083, "updates_nfe_3": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.508511284310853, "grad_norm": 0.04289498554100681, "grad_rms": 0.0001249060080945649, "grad_to_weight": 0.0025983557694734595, "update_to_weight": 4.6298991634650515e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25827238357254, "grad_norm": 0.017380762950478856, "grad_rms": 5.061108403234748e-05, "grad_to_weight": 0.001069041195793994, "update_to_weight": 5.61015701509889e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065648171467767, "grad_norm": 0.0012637042070487835, "grad_rms": 2.2623934476346025e-05, "grad_to_weight": 0.000249465451265761, "update_to_weight": 2.8498137691062474e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.865056911302595, "grad_norm": 0.012644661561212267, "grad_rms": 4.046296878851245e-05, "grad_to_weight": 0.0009828686844053757, "update_to_weight": 5.920356618696278e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043998536293489, "grad_norm": 0.018961358648824807, "grad_rms": 0.00043138584625429207, "grad_to_weight": 0.0037591919411535542, "update_to_weight": 2.1647877388829035e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.643291306343603, "grad_norm": 0.07552099296399488, "grad_rms": 0.00018197882461664836, "grad_to_weight": 0.004537623693169992, "update_to_weight": 6.101434749752432e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.572346895109423, "grad_norm": 0.023394728167153062, "grad_rms": 3.227654363588287e-05, "grad_to_weight": 0.0009520754475352642, "update_to_weight": 7.285139032033642e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831902328138543, "grad_norm": 0.010527086499734457, "grad_rms": 1.2779637767845856e-05, "grad_to_weight": 0.0003414348679395944, "update_to_weight": 5.897195236705318e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63041650154258, "grad_norm": 0.0068856271012882305, "grad_rms": 6.181242855903021e-06, "grad_to_weight": 0.0001988317726695986, "update_to_weight": 7.418994763038665e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08072742921432, "grad_norm": 0.005643008828684285, "grad_rms": 4.539020824120743e-06, "grad_to_weight": 0.00015218171864229393, "update_to_weight": 7.570314906661127e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670488353607205, "grad_norm": 0.0019384312580302682, "grad_rms": 3.502147021189413e-06, "grad_to_weight": 7.857287744962048e-05, "update_to_weight": 5.1339486246074435e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.351997977747956, "grad_norm": 0.00582350333957013, "grad_rms": 4.684203716514225e-06, "grad_to_weight": 0.0001559087506654776, "update_to_weight": 7.627521271661906e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.492377356828236, "grad_norm": 0.014733428560796206, "grad_rms": 6.541131795822771e-06, "grad_to_weight": 0.00023205035272177773, "update_to_weight": 8.496761860582703e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.77993988694625, "grad_norm": 0.025768460457351563, "grad_rms": 1.9227678104332025e-05, "grad_to_weight": 0.0005074535439530076, "update_to_weight": 7.241138476449584e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10290717719228, "grad_norm": 0.03337617614656953, "grad_rms": 2.8662142818788128e-05, "grad_to_weight": 0.0008322632571025065, "update_to_weight": 8.680948751853757e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20358379128218, "grad_norm": 0.06832687763516636, "grad_rms": 0.00010651191971751743, "grad_to_weight": 0.002422631043657878, "update_to_weight": 7.222225201928648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414689942947245, "grad_norm": 0.009551573856498493, "grad_rms": 0.0008130839624368868, "grad_to_weight": 0.0011351070474681156, "update_to_weight": 5.966135158304345e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3202689349870478, "grad_norm": 0.12686751486867984, "grad_rms": 0.0029369368816895337, "grad_to_weight": 0.09609217600043315, "update_to_weight": 0.0001835085272722861, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47650}
{"step": 47700, "lr": 1.0520103669039554e-05, "window_steps": 50, "elapsed_seconds": 3862.123675582, "steps_per_second": 1.7347968534410925, "eta_seconds": 1325.8036498266567, "images_seen": 48844800, "images_per_second": 1776.4319779236787, "effective_batch_size": 1024, "loss": 0.3843117952346802, "edge_loss": 0.20497357845306396, "diag_loss": 0.17933814227581024, "jvp_rms": 2.128164052963257, "grad_norm": 0.4613628387451172, "grad_norm_postclip": 0.4531759023666382, "clip_fraction": 0.07999999821186066, "ema_updates": 47700, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4351392765839895, "updates_nfe_1": 15, "loss_nfe_5": 0.3707083788785068, "updates_nfe_5": 11, "loss_nfe_4": 0.35859269174662506, "updates_nfe_4": 11, "loss_nfe_3": 0.35245231539011, "updates_nfe_3": 10, "loss_nfe_2": 0.3805546164512634, "updates_nfe_2": 3, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.505958174957232, "grad_norm": 0.031526087297394687, "grad_rms": 9.180088687508424e-05, "grad_to_weight": 0.0019099822599348353, "update_to_weight": 3.7666661786465405e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26088408784588, "grad_norm": 0.015141746816458157, "grad_rms": 4.409128774770963e-05, "grad_to_weight": 0.0009311761116221094, "update_to_weight": 6.000230751082124e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066093770481217, "grad_norm": 0.001110197182411819, "grad_rms": 1.9875717886044327e-05, "grad_to_weight": 0.00021914264376246708, "update_to_weight": 2.3174167804005234e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.8667279678492, "grad_norm": 0.013934488124298927, "grad_rms": 4.4590419073529505e-05, "grad_to_weight": 0.0010829861452824523, "update_to_weight": 6.620039175762022e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043749712597769, "grad_norm": 0.033688666052578004, "grad_rms": 0.0007664436912684168, "grad_to_weight": 0.006679289808618745, "update_to_weight": 2.704119741543622e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.643669945885676, "grad_norm": 0.113562547079065, "grad_rms": 0.0002736454861997601, "grad_to_weight": 0.006823167453349898, "update_to_weight": 8.286182381054973e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.572811379287298, "grad_norm": 0.021532162708224296, "grad_rms": 2.970685464953214e-05, "grad_to_weight": 0.0008762596341082079, "update_to_weight": 0.0001101540699101926, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831916185408627, "grad_norm": 0.0085601764141759, "grad_rms": 1.0391854745829734e-05, "grad_to_weight": 0.0002776401039331785, "update_to_weight": 8.808381179726079e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.630039330376704, "grad_norm": 0.004567069134445644, "grad_rms": 4.099868181131354e-06, "grad_to_weight": 0.00013188171953473676, "update_to_weight": 9.225276866926577e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08019741100447, "grad_norm": 0.0035793199958799194, "grad_rms": 2.879068328744536e-06, "grad_to_weight": 9.652915156319171e-05, "update_to_weight": 8.545021093319003e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671042526724904, "grad_norm": 0.0012705030570277494, "grad_rms": 2.295406906047889e-06, "grad_to_weight": 5.149774500414715e-05, "update_to_weight": 6.37709633461596e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35217282104382, "grad_norm": 0.0038223876919699345, "grad_rms": 3.074582701966042e-06, "grad_to_weight": 0.00010233374401760214, "update_to_weight": 8.208049391903416e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.492383124621725, "grad_norm": 0.010139657284977002, "grad_rms": 4.5016565147633775e-06, "grad_to_weight": 0.00015969879828065458, "update_to_weight": 9.158759016748723e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.780170661606434, "grad_norm": 0.029434934797852936, "grad_rms": 2.1963494957404475e-05, "grad_to_weight": 0.0005796541132956043, "update_to_weight": 7.196079737710654e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10357074773704, "grad_norm": 0.0442647319100029, "grad_rms": 3.8012804770335355e-05, "grad_to_weight": 0.0011037603655904048, "update_to_weight": 8.998478142513194e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20395601268774, "grad_norm": 0.11177056218168012, "grad_rms": 0.000174234467575754, "grad_to_weight": 0.003962939175319922, "update_to_weight": 7.3721611501022e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414135492516493, "grad_norm": 0.018248436012587695, "grad_rms": 0.0015534100331848412, "grad_to_weight": 0.0021687832373055793, "update_to_weight": 6.031765057815891e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3201645988374637, "grad_norm": 0.24919164623813256, "grad_rms": 0.005768696085859714, "grad_to_weight": 0.1887580127944429, "update_to_weight": 0.0002168698783235893, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47700}
{"step": 47750, "lr": 1.0497787189308258e-05, "window_steps": 50, "elapsed_seconds": 3889.634414707, "steps_per_second": 1.7353816015401713, "eta_seconds": 1296.5448049023332, "images_seen": 48896000, "images_per_second": 1777.0307599771354, "effective_batch_size": 1024, "loss": 0.38720703125, "edge_loss": 0.20064853131771088, "diag_loss": 0.18655845522880554, "jvp_rms": 2.565234422683716, "grad_norm": 0.3104249835014343, "grad_norm_postclip": 0.3104249835014343, "clip_fraction": 0.0, "ema_updates": 47750, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.347938422113657, "updates_nfe_2": 8, "loss_nfe_1": 0.4003053108851115, "updates_nfe_1": 12, "loss_nfe_3": 0.39553987979888916, "updates_nfe_3": 10, "loss_nfe_4": 0.3912010043859482, "updates_nfe_4": 10, "loss_nfe_5": 0.3905769720673561, "updates_nfe_5": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.509208190525825, "grad_norm": 0.2862714490184063, "grad_rms": 0.000833594497756686, "grad_to_weight": 0.017340107757723322, "update_to_weight": 6.925503769317586e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.256590460439035, "grad_norm": 0.11599891217904527, "grad_rms": 0.00033777750198202945, "grad_to_weight": 0.007135500673485782, "update_to_weight": 4.989291295319254e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066482484548363, "grad_norm": 0.00731505426606007, "grad_rms": 0.0001309604791082806, "grad_to_weight": 0.0014438131955196427, "update_to_weight": 2.951066833487797e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.865687882216683, "grad_norm": 0.09751721498384748, "grad_rms": 0.0003120554873795914, "grad_to_weight": 0.00757963475226525, "update_to_weight": 7.584867737422593e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.044551917026833, "grad_norm": 0.03876665258410241, "grad_rms": 0.0008819718850935634, "grad_to_weight": 0.007684855507830865, "update_to_weight": 2.898544266499572e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.643999987854578, "grad_norm": 0.23645531565607358, "grad_rms": 0.0005697734991112419, "grad_to_weight": 0.014206639980090077, "update_to_weight": 6.649121749438952e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57260770281425, "grad_norm": 0.037749795440745644, "grad_rms": 5.208151644606829e-05, "grad_to_weight": 0.0015362551625492413, "update_to_weight": 6.090285057598727e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832054757766915, "grad_norm": 0.0138513275616105, "grad_rms": 1.681518897423584e-05, "grad_to_weight": 0.0004492508744692472, "update_to_weight": 5.1039904770649607e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63149864274123, "grad_norm": 0.003867633450613825, "grad_rms": 3.471983202718826e-06, "grad_to_weight": 0.00011167964431780321, "update_to_weight": 5.3484597580053364e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.081104363020934, "grad_norm": 0.002729485285641273, "grad_rms": 2.195493738673724e-06, "grad_to_weight": 7.360852198251268e-05, "update_to_weight": 5.7552359204638704e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67078770759623, "grad_norm": 0.001144505685472427, "grad_rms": 2.067768542478291e-06, "grad_to_weight": 4.639112861078324e-05, "update_to_weight": 3.7214813988646994e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35226596153087, "grad_norm": 0.0030006310244315576, "grad_rms": 2.413593017286347e-06, "grad_to_weight": 8.033330635206738e-05, "update_to_weight": 5.814563552394345e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49374430923294, "grad_norm": 0.009288418505961241, "grad_rms": 4.123735990678898e-06, "grad_to_weight": 0.00014628871878659338, "update_to_weight": 6.294827798980718e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78012979533678, "grad_norm": 0.03468567772000788, "grad_rms": 2.588144709439339e-05, "grad_to_weight": 0.0006830561059966633, "update_to_weight": 5.493380908321609e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10381273518065, "grad_norm": 0.10838020016641492, "grad_rms": 9.30726384669417e-05, "grad_to_weight": 0.0027024911791326896, "update_to_weight": 7.917733159140111e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203730949169206, "grad_norm": 0.2307119218109619, "grad_rms": 0.0003596471922076511, "grad_to_weight": 0.008180191557874649, "update_to_weight": 7.05078208782622e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414844983056264, "grad_norm": 0.003262759291607236, "grad_rms": 0.0002777445155274448, "grad_to_weight": 0.0003877384905101608, "update_to_weight": 5.620032065202864e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3206619612815205, "grad_norm": 0.03908181458040921, "grad_rms": 0.0009047298101753222, "grad_to_weight": 0.029592595021428268, "update_to_weight": 0.00011761872926874707, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47750}
{"step": 47800, "lr": 1.047595737864637e-05, "window_steps": 50, "elapsed_seconds": 3917.143611819, "steps_per_second": 1.735958819452701, "eta_seconds": 1267.3111685296765, "images_seen": 48947200, "images_per_second": 1777.6218311195657, "effective_batch_size": 1024, "loss": 0.3710395097732544, "edge_loss": 0.19241361320018768, "diag_loss": 0.1786258965730667, "jvp_rms": 2.1532421112060547, "grad_norm": 0.20597554743289948, "grad_norm_postclip": 0.20597554743289948, "clip_fraction": 0.0, "ema_updates": 47800, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.361712720990181, "updates_nfe_4": 10, "loss_nfe_1": 0.4032169580459595, "updates_nfe_1": 11, "loss_nfe_5": 0.3579292916334592, "updates_nfe_5": 13, "loss_nfe_3": 0.3622048483653502, "updates_nfe_3": 11, "loss_nfe_2": 0.3724256753921509, "updates_nfe_2": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.508532543158267, "grad_norm": 0.0650145792486243, "grad_rms": 0.00018931610442274538, "grad_to_weight": 0.003938240971973532, "update_to_weight": 4.220373356221875e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.257892276654342, "grad_norm": 0.06527275509708722, "grad_rms": 0.00019006788727594397, "grad_to_weight": 0.004014835009751921, "update_to_weight": 3.800453728204483e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0659145566630075, "grad_norm": 0.0035935966799776863, "grad_rms": 6.433570084576863e-05, "grad_to_weight": 0.0007093678031444813, "update_to_weight": 2.487385037426814e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866621828412873, "grad_norm": 0.02491269060299495, "grad_rms": 7.972071197216047e-05, "grad_to_weight": 0.0019362262243521606, "update_to_weight": 5.167931125935366e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.044025006177057, "grad_norm": 0.01967808395387595, "grad_rms": 0.0004476919115514914, "grad_to_weight": 0.003901266137613832, "update_to_weight": 1.6339229378698287e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64430527084921, "grad_norm": 0.07744582417358609, "grad_rms": 0.00018661698557507017, "grad_to_weight": 0.004652992294561221, "update_to_weight": 4.758393089433408e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.573360303814653, "grad_norm": 0.014455728723252199, "grad_rms": 1.994385040898357e-05, "grad_to_weight": 0.0005882682931649425, "update_to_weight": 6.587911140507322e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832107217140038, "grad_norm": 0.006698638540126269, "grad_rms": 8.131991133796126e-06, "grad_to_weight": 0.00021726178145885512, "update_to_weight": 6.384386747371266e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63031427798754, "grad_norm": 0.0030088763227538065, "grad_rms": 2.701075007509209e-06, "grad_to_weight": 8.688561988206884e-05, "update_to_weight": 8.170838505326965e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08043608473847, "grad_norm": 0.0025491495056578044, "grad_rms": 2.0504385233569917e-06, "grad_to_weight": 6.874648129359461e-05, "update_to_weight": 8.103298079135531e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670380733723487, "grad_norm": 0.0009845790175608542, "grad_rms": 1.7788304120622586e-06, "grad_to_weight": 3.990935641357864e-05, "update_to_weight": 5.246667560534815e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.352383612340425, "grad_norm": 0.002684485296110141, "grad_rms": 2.1592974654145437e-06, "grad_to_weight": 7.186918307465778e-05, "update_to_weight": 7.781176268343426e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49335402905862, "grad_norm": 0.007177791840540581, "grad_rms": 3.1866908804164982e-06, "grad_to_weight": 0.00011304792368120236, "update_to_weight": 8.430949501518852e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78035335746863, "grad_norm": 0.017676321734020995, "grad_rms": 1.3189558799326713e-05, "grad_to_weight": 0.0003480937127315443, "update_to_weight": 5.8564966175552646e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10409733489678, "grad_norm": 0.033875487717323405, "grad_rms": 2.9090931889446602e-05, "grad_to_weight": 0.0008446889462301021, "update_to_weight": 5.953292391481899e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204083692154626, "grad_norm": 0.0907635819235878, "grad_rms": 0.00014148756222608213, "grad_to_weight": 0.0032181007159908194, "update_to_weight": 4.5311874327764845e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414499538905114, "grad_norm": 0.005065735381351924, "grad_rms": 0.0004312240326471485, "grad_to_weight": 0.0006020245598600476, "update_to_weight": 3.0513621676289954e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3203032002319022, "grad_norm": 0.06067099654109187, "grad_rms": 0.0014045115298020262, "grad_to_weight": 0.04595232105052493, "update_to_weight": 4.8971664329317024e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47800}
{"step": 47850, "lr": 1.0454614475781217e-05, "window_steps": 50, "elapsed_seconds": 3944.673904451, "steps_per_second": 1.736518699878019, "eta_seconds": 1238.1093276744016, "images_seen": 48998400, "images_per_second": 1778.1951486750916, "effective_batch_size": 1024, "loss": 0.35808536410331726, "edge_loss": 0.1842300146818161, "diag_loss": 0.1738552749156952, "jvp_rms": 1.9960546493530273, "grad_norm": 0.1713005155324936, "grad_norm_postclip": 0.1713005155324936, "clip_fraction": 0.0, "ema_updates": 47850, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3516547933220863, "updates_nfe_5": 10, "loss_nfe_2": 0.3736545497720892, "updates_nfe_2": 11, "loss_nfe_3": 0.33319914882833307, "updates_nfe_3": 11, "loss_nfe_1": 0.39744503423571587, "updates_nfe_1": 8, "loss_nfe_4": 0.34327669739723204, "updates_nfe_4": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.508254327381195, "grad_norm": 0.4153245370468143, "grad_rms": 0.0012093844847353501, "grad_to_weight": 0.025158598166127224, "update_to_weight": 5.4331770019896965e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.258210440829004, "grad_norm": 0.3196419582388554, "grad_rms": 0.0009307661611163687, "grad_to_weight": 0.01966034081070468, "update_to_weight": 4.888411428209308e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066242671258173, "grad_norm": 0.011640337280023524, "grad_rms": 0.000208395466626514, "grad_to_weight": 0.0022976272625197225, "update_to_weight": 2.8261211278975813e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866169391832829, "grad_norm": 0.14198363486955873, "grad_rms": 0.00045434821314867295, "grad_to_weight": 0.011035424029135425, "update_to_weight": 5.536208162330581e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043977360286621, "grad_norm": 0.03583436542083214, "grad_rms": 0.0008152600421916128, "grad_to_weight": 0.007104386649902781, "update_to_weight": 1.594014506242557e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64386797185237, "grad_norm": 0.25435854900785343, "grad_rms": 0.0006129139456854514, "grad_to_weight": 0.015282418091636949, "update_to_weight": 4.486530218721782e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57292066840613, "grad_norm": 0.030747257913311277, "grad_rms": 4.242046347501993e-05, "grad_to_weight": 0.0012512659088524064, "update_to_weight": 4.645744508298456e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832279441624266, "grad_norm": 0.012621879829836195, "grad_rms": 1.5322668069522876e-05, "grad_to_weight": 0.0004093722572063996, "update_to_weight": 4.0267150154579145e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63117083155091, "grad_norm": 0.003522483843439284, "grad_rms": 3.162141628061636e-06, "grad_to_weight": 0.00010171425796063778, "update_to_weight": 4.8049610000164706e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08136278267443, "grad_norm": 0.0026258315258197522, "grad_rms": 2.112118612280737e-06, "grad_to_weight": 7.081270289900517e-05, "update_to_weight": 5.33074625162289e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67025084703171, "grad_norm": 0.0011353305485065225, "grad_rms": 2.0511919017225086e-06, "grad_to_weight": 4.6020227177508565e-05, "update_to_weight": 3.4734627068586655e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.353922842960046, "grad_norm": 0.0026900490937125654, "grad_rms": 2.163772771752933e-06, "grad_to_weight": 7.201516973255591e-05, "update_to_weight": 5.6690944022511916e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49463060160579, "grad_norm": 0.008229599551068292, "grad_rms": 3.653657060761731e-06, "grad_to_weight": 0.00012961095250249655, "update_to_weight": 6.633545032769487e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78059855351377, "grad_norm": 0.03831717553408241, "grad_rms": 2.8591165477498808e-05, "grad_to_weight": 0.0007545632904209053, "update_to_weight": 5.176020558976381e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10450064119596, "grad_norm": 0.1713035468629242, "grad_rms": 0.00014710872521730596, "grad_to_weight": 0.004271429493550621, "update_to_weight": 6.478583681181649e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203846727261226, "grad_norm": 0.3817453496458628, "grad_rms": 0.0005950869034455735, "grad_to_weight": 0.013535222813307804, "update_to_weight": 5.238289693988482e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41494335497753, "grad_norm": 0.005537116194313802, "grad_rms": 0.0004713506321979652, "grad_to_weight": 0.0006580099188711161, "update_to_weight": 2.0041843880978e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3207482967419881, "grad_norm": 0.103229374337621, "grad_rms": 0.002389722515489587, "grad_to_weight": 0.07815976336465202, "update_to_weight": 5.424415154514125e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47850}
{"step": 47900, "lr": 1.043375871411535e-05, "window_steps": 50, "elapsed_seconds": 3972.193676867, "steps_per_second": 1.7370754201094891, "eta_seconds": 1208.928510350826, "images_seen": 49049600, "images_per_second": 1778.7652301921169, "effective_batch_size": 1024, "loss": 0.3612951934337616, "edge_loss": 0.1869155913591385, "diag_loss": 0.17437949776649475, "jvp_rms": 1.9868358373641968, "grad_norm": 0.16473594307899475, "grad_norm_postclip": 0.16473594307899475, "clip_fraction": 0.0, "ema_updates": 47900, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3515124171972275, "updates_nfe_2": 10, "loss_nfe_3": 0.3678460137711631, "updates_nfe_3": 9, "loss_nfe_5": 0.37717320973222906, "updates_nfe_5": 11, "loss_nfe_1": 0.3994982659816742, "updates_nfe_1": 10, "loss_nfe_4": 0.3095131069421768, "updates_nfe_4": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50743351730138, "grad_norm": 0.013437357077178108, "grad_rms": 3.912827130451046e-05, "grad_to_weight": 0.0008140185488613033, "update_to_weight": 3.122183534489844e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25886833578976, "grad_norm": 0.0041699897769797685, "grad_rms": 1.2142602923592513e-05, "grad_to_weight": 0.0002564747859972885, "update_to_weight": 4.015909891860498e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06547440170283, "grad_norm": 0.0006676499074391763, "grad_rms": 1.195285073420612e-05, "grad_to_weight": 0.00013180402357077086, "update_to_weight": 2.096860269225947e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866468250806957, "grad_norm": 0.0062778577325236515, "grad_rms": 2.008917045823033e-05, "grad_to_weight": 0.00048792392831886234, "update_to_weight": 4.6214633516795855e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043700929613374, "grad_norm": 0.012533866928947963, "grad_rms": 0.000285155346308373, "grad_to_weight": 0.002485053555685101, "update_to_weight": 1.7955677945728032e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644156755497182, "grad_norm": 0.06378291735859745, "grad_rms": 0.00015369422297019004, "grad_to_weight": 0.0038321507238587753, "update_to_weight": 5.469412846357927e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57320879193499, "grad_norm": 0.016571474333199582, "grad_rms": 2.2862839465576847e-05, "grad_to_weight": 0.0006743716082629955, "update_to_weight": 7.193653759170157e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832015165728098, "grad_norm": 0.007752288470477934, "grad_rms": 9.411097603031462e-06, "grad_to_weight": 0.00025143632126566724, "update_to_weight": 6.522484597548375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63060508558506, "grad_norm": 0.004901564992050184, "grad_rms": 4.400145860961036e-06, "grad_to_weight": 0.00014153853159471513, "update_to_weight": 8.140030759501959e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08107473517488, "grad_norm": 0.004084049659764876, "grad_rms": 3.2850535973267652e-06, "grad_to_weight": 0.00011013838430876901, "update_to_weight": 8.359219545039013e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670560099935535, "grad_norm": 0.0013910395870201395, "grad_rms": 2.5131792143041595e-06, "grad_to_weight": 5.6384596919782714e-05, "update_to_weight": 5.713629647394423e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.353728400129285, "grad_norm": 0.004184109965444948, "grad_rms": 3.365538286424023e-06, "grad_to_weight": 0.00011201318167292951, "update_to_weight": 8.414619900130128e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49496704341022, "grad_norm": 0.010978863418648142, "grad_rms": 4.874234961223035e-06, "grad_to_weight": 0.00017290919154493167, "update_to_weight": 8.867881708972197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78071874722112, "grad_norm": 0.022464742804648555, "grad_rms": 1.676253977451573e-05, "grad_to_weight": 0.00044238725561319265, "update_to_weight": 6.405850569237213e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.104693922399846, "grad_norm": 0.026002076846558857, "grad_rms": 2.2329557373149678e-05, "grad_to_weight": 0.000648354950591851, "update_to_weight": 6.561430815752078e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20401877046403, "grad_norm": 0.045379389830819655, "grad_rms": 7.074003809011346e-05, "grad_to_weight": 0.001608968927447393, "update_to_weight": 4.940051763295951e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414778343359442, "grad_norm": 0.004082519240277237, "grad_rms": 0.0003475271165234243, "grad_to_weight": 0.00048516063925783313, "update_to_weight": 1.7940133809858624e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.320637318760854, "grad_norm": 0.06181366524186249, "grad_rms": 0.0014309638951243804, "grad_to_weight": 0.046805935561371144, "update_to_weight": 5.21130255722456e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47900}
{"step": 47950, "lr": 1.0413390321724046e-05, "window_steps": 50, "elapsed_seconds": 3999.695559745, "steps_per_second": 1.7376322513014206, "eta_seconds": 1179.766316183777, "images_seen": 49100800, "images_per_second": 1779.3354253326547, "effective_batch_size": 1024, "loss": 0.3437623977661133, "edge_loss": 0.17875269055366516, "diag_loss": 0.16500964760780334, "jvp_rms": 1.67578125, "grad_norm": 0.18141904473304749, "grad_norm_postclip": 0.18141904473304749, "clip_fraction": 0.0, "ema_updates": 47950, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3041523424061862, "updates_nfe_3": 11, "loss_nfe_4": 0.326797821692058, "updates_nfe_4": 7, "loss_nfe_1": 0.3986399521430333, "updates_nfe_1": 12, "loss_nfe_5": 0.3325610106641596, "updates_nfe_5": 11, "loss_nfe_2": 0.3458896213107639, "updates_nfe_2": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.510070501924442, "grad_norm": 0.15279756102450368, "grad_rms": 0.00044493157308355324, "grad_to_weight": 0.009254809723961708, "update_to_weight": 6.441173136265347e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25653789753718, "grad_norm": 0.07618757232900414, "grad_rms": 0.00022185076894208177, "grad_to_weight": 0.004686580427468898, "update_to_weight": 8.034917833260771e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065998140367121, "grad_norm": 0.004516219351680065, "grad_rms": 8.085329630407048e-05, "grad_to_weight": 0.0008914767093366492, "update_to_weight": 2.6073429931014773e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.865954730904944, "grad_norm": 0.03626014851050255, "grad_rms": 0.0001160326237554616, "grad_to_weight": 0.0028183021990123343, "update_to_weight": 6.317333757811074e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043846331767201, "grad_norm": 0.01960902004400371, "grad_rms": 0.0004461206531961355, "grad_to_weight": 0.003887711630011802, "update_to_weight": 1.8609380026889505e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64427043403375, "grad_norm": 0.09221152744106115, "grad_rms": 0.00022219709674408487, "grad_to_weight": 0.005540136337397495, "update_to_weight": 5.3120580363927685e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57313924518561, "grad_norm": 0.03519715307241113, "grad_rms": 4.855976264753382e-05, "grad_to_weight": 0.0014323425558786503, "update_to_weight": 7.486076106571239e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831991410480406, "grad_norm": 0.024649624465332157, "grad_rms": 2.9924069854306052e-05, "grad_to_weight": 0.000799482074873479, "update_to_weight": 6.849441218256816e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.6311250082341, "grad_norm": 0.017386636889262206, "grad_rms": 1.5608022839317745e-05, "grad_to_weight": 0.0005020523267762239, "update_to_weight": 7.954139015918559e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.081171848581896, "grad_norm": 0.01414275590390149, "grad_rms": 1.1375892809516103e-05, "grad_to_weight": 0.0003813999180406796, "update_to_weight": 8.358665897510876e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67067019303362, "grad_norm": 0.005238677206369428, "grad_rms": 9.464672888066424e-06, "grad_to_weight": 0.0002123443410892299, "update_to_weight": 5.455944768265669e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.354055194223896, "grad_norm": 0.013522847849797731, "grad_rms": 1.0877262441916062e-05, "grad_to_weight": 0.00036201820068758653, "update_to_weight": 8.403901954069034e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49538807378553, "grad_norm": 0.02656361037626241, "grad_rms": 1.1793322628675903e-05, "grad_to_weight": 0.0004183549574560261, "update_to_weight": 9.115853147255879e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78083894064399, "grad_norm": 0.03251968539834361, "grad_rms": 2.426524642123545e-05, "grad_to_weight": 0.0006403928347137938, "update_to_weight": 6.733113996053719e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.104920684199556, "grad_norm": 0.04378492438732855, "grad_rms": 3.760076500640228e-05, "grad_to_weight": 0.0010917594060865163, "update_to_weight": 7.176024184875237e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20395709471955, "grad_norm": 0.10799809148269948, "grad_rms": 0.0001683537203481109, "grad_to_weight": 0.0038291822356700186, "update_to_weight": 5.311391906566269e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.415189053094197, "grad_norm": 0.0029968198350812185, "grad_rms": 0.0002551062456120243, "grad_to_weight": 0.00035612032197652316, "update_to_weight": 2.6514511373851883e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3209853404731156, "grad_norm": 0.051368309414871645, "grad_rms": 0.0011891577022434515, "grad_to_weight": 0.03888635841823491, "update_to_weight": 7.493367499193247e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 47950}
{"step": 48000, "lr": 1.0393509521352768e-05, "window_steps": 50, "elapsed_seconds": 4027.1706882820004, "steps_per_second": 1.7381930247873887, "eta_seconds": 1150.620196652, "images_seen": 49152000, "images_per_second": 1779.909657382286, "effective_batch_size": 1024, "loss": 0.38119250535964966, "edge_loss": 0.19608590006828308, "diag_loss": 0.18510659039020538, "jvp_rms": 2.4033594131469727, "grad_norm": 0.3364473283290863, "grad_norm_postclip": 0.3241864740848541, "clip_fraction": 0.019999999552965164, "ema_updates": 48000, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4035877101123333, "updates_nfe_1": 8, "loss_nfe_5": 0.43281155493524337, "updates_nfe_5": 9, "loss_nfe_3": 0.30599088470141095, "updates_nfe_3": 12, "loss_nfe_4": 0.40004164251414215, "updates_nfe_4": 11, "loss_nfe_2": 0.38632703125476836, "updates_nfe_2": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.505217680883806, "grad_norm": 0.049463894279366975, "grad_rms": 0.00014403402871743454, "grad_to_weight": 0.002996864096900437, "update_to_weight": 8.83301422431725e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.261937847581198, "grad_norm": 0.09634607727872704, "grad_rms": 0.0002805503663581325, "grad_to_weight": 0.005924636914846995, "update_to_weight": 8.077311300160739e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065751526821479, "grad_norm": 0.0038194488496447183, "grad_rms": 6.837910329658461e-05, "grad_to_weight": 0.0007539747714474348, "update_to_weight": 4.097896463313001e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866572019660458, "grad_norm": 0.0515866824360173, "grad_rms": 0.0001650775950947123, "grad_to_weight": 0.0040093571432384244, "update_to_weight": 9.833884819650977e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043115875201847, "grad_norm": 0.025545367397291012, "grad_rms": 0.000581177231898421, "grad_to_weight": 0.005065393702909628, "update_to_weight": 3.432360200247364e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64371211827974, "grad_norm": 0.10407428867277747, "grad_rms": 0.000250782146555134, "grad_to_weight": 0.006253069503555819, "update_to_weight": 8.642724343821306e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.572862298141775, "grad_norm": 0.02376913469410074, "grad_rms": 3.279309371153424e-05, "grad_to_weight": 0.0009672920641360606, "update_to_weight": 9.609914920641255e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83210820693867, "grad_norm": 0.009905922317035085, "grad_rms": 1.20255589209153e-05, "grad_to_weight": 0.0003212859221480608, "update_to_weight": 8.83667065059334e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.631655497535775, "grad_norm": 0.003521114487769247, "grad_rms": 3.160912354412614e-06, "grad_to_weight": 0.00010167329390360195, "update_to_weight": 0.0001044406436371506, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08184669743657, "grad_norm": 0.0029481203017345903, "grad_rms": 2.3713553970650904e-06, "grad_to_weight": 7.95030605079976e-05, "update_to_weight": 9.664029894471662e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.6708879037812, "grad_norm": 0.001112394355410649, "grad_rms": 2.0097532796432498e-06, "grad_to_weight": 4.508935226608351e-05, "update_to_weight": 6.779567319520386e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.355221825820415, "grad_norm": 0.003141327488212144, "grad_rms": 2.5267638802724206e-06, "grad_to_weight": 8.409339671062581e-05, "update_to_weight": 8.571420124729803e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49639545933395, "grad_norm": 0.008666470841053422, "grad_rms": 3.847612776759e-06, "grad_to_weight": 0.00013648760340425676, "update_to_weight": 7.77345213854476e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78108173049044, "grad_norm": 0.02334999847167687, "grad_rms": 1.7423091887585417e-05, "grad_to_weight": 0.00045981687817526053, "update_to_weight": 5.955737657820327e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10488415889563, "grad_norm": 0.05842102939045259, "grad_rms": 5.0169674340656205e-05, "grad_to_weight": 0.0014567060999101347, "update_to_weight": 8.354642229078236e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203749343851097, "grad_norm": 0.15912041994202455, "grad_rms": 0.00024804618593546996, "grad_to_weight": 0.005641817972570925, "update_to_weight": 7.52454582711471e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.414927941941416, "grad_norm": 0.005462529271748981, "grad_rms": 0.0004650013717037151, "grad_to_weight": 0.000649147480458248, "update_to_weight": 2.8823554345885567e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3208782631646148, "grad_norm": 0.07685198442735534, "grad_rms": 0.0017790955212558606, "grad_to_weight": 0.0581824885536614, "update_to_weight": 7.551612442269815e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48000}
{"step": 48050, "lr": 1.0374116530414776e-05, "window_steps": 50, "elapsed_seconds": 4077.9066379210003, "steps_per_second": 1.728828201813427, "eta_seconds": 1127.9316232547446, "images_seen": 49203200, "images_per_second": 1770.3200786569491, "effective_batch_size": 1024, "loss": 0.3875241279602051, "edge_loss": 0.19996421039104462, "diag_loss": 0.18755991756916046, "jvp_rms": 2.5307812690734863, "grad_norm": 0.2809384763240814, "grad_norm_postclip": 0.2809384763240814, "clip_fraction": 0.0, "ema_updates": 48050, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.4143030991156896, "updates_nfe_5": 12, "loss_nfe_4": 0.37399367414988005, "updates_nfe_4": 13, "loss_nfe_1": 0.39975821375846865, "updates_nfe_1": 10, "loss_nfe_3": 0.3545506106955664, "updates_nfe_3": 7, "loss_nfe_2": 0.38290200755000114, "updates_nfe_2": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.509090809178534, "grad_norm": 0.0967389512727217, "grad_rms": 0.0002816943770543741, "grad_to_weight": 0.005859738273348032, "update_to_weight": 3.469115281306483e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25597846755156, "grad_norm": 0.034170178384900135, "grad_rms": 9.950022185826072e-05, "grad_to_weight": 0.0021020068680028692, "update_to_weight": 4.604809312806197e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065610895249005, "grad_norm": 0.004443358286068857, "grad_rms": 7.954887398350675e-05, "grad_to_weight": 0.0008771613884193605, "update_to_weight": 2.404957266001931e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.864954316302367, "grad_norm": 0.034019180197376764, "grad_rms": 0.00010886151597443527, "grad_to_weight": 0.00264432965411062, "update_to_weight": 5.2478562983069724e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043502767523344, "grad_norm": 0.06717226748452763, "grad_rms": 0.0015282220008759012, "grad_to_weight": 0.013318574526630659, "update_to_weight": 2.3104424551257595e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64396148329543, "grad_norm": 0.13077955516049689, "grad_rms": 0.00031513237310507465, "grad_to_weight": 0.007857477637866001, "update_to_weight": 5.5181317013057676e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.572601493139164, "grad_norm": 0.04642004020895083, "grad_rms": 6.404342220514745e-05, "grad_to_weight": 0.0018890975065017687, "update_to_weight": 7.315891487327961e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832208176437124, "grad_norm": 0.03072424528616615, "grad_rms": 3.729851801422486e-05, "grad_to_weight": 0.000996498373076195, "update_to_weight": 6.242424741552022e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63251201792725, "grad_norm": 0.02182173689218163, "grad_rms": 1.9589422035787765e-05, "grad_to_weight": 0.0006300939672203327, "update_to_weight": 7.501751695792957e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08266966722376, "grad_norm": 0.016263641362317956, "grad_rms": 1.3081852086487811e-05, "grad_to_weight": 0.0004385779532128155, "update_to_weight": 7.851361462808718e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671169935302267, "grad_norm": 0.006020912961078115, "grad_rms": 1.0877931473776923e-05, "grad_to_weight": 0.0002440465116517527, "update_to_weight": 5.249688953637411e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35515810318653, "grad_norm": 0.014971329937339952, "grad_rms": 1.2042366123005458e-05, "grad_to_weight": 0.000400783471347772, "update_to_weight": 8.292574924819177e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49671266712356, "grad_norm": 0.029803602572601923, "grad_rms": 1.3231766904306797e-05, "grad_to_weight": 0.0004693723709579254, "update_to_weight": 9.399027926800385e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7814038459208, "grad_norm": 0.04456762649690518, "grad_rms": 3.325507077052074e-05, "grad_to_weight": 0.0008776367552210796, "update_to_weight": 7.159546952213083e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.105170272886106, "grad_norm": 0.04959100576544003, "grad_rms": 4.2586798545599134e-05, "grad_to_weight": 0.0012365240049602036, "update_to_weight": 7.291494732528216e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203987391593344, "grad_norm": 0.13213995036611192, "grad_rms": 0.0002059874572349582, "grad_to_weight": 0.004685151377053103, "update_to_weight": 5.843128769152775e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41539530704459, "grad_norm": 0.009985505119789205, "grad_rms": 0.0008500226446145481, "grad_to_weight": 0.001186575883301675, "update_to_weight": 2.4827660706246814e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3212639805569282, "grad_norm": 0.10341646285216935, "grad_rms": 0.0023940535466372083, "grad_to_weight": 0.0782708560696388, "update_to_weight": 7.649333801606112e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48050}
{"step": 48100, "lr": 1.0355211560988715e-05, "window_steps": 50, "elapsed_seconds": 4105.373787314, "steps_per_second": 1.7294405741907555, "eta_seconds": 1098.6211543516338, "images_seen": 49254400, "images_per_second": 1770.9471479713336, "effective_batch_size": 1024, "loss": 0.36434119939804077, "edge_loss": 0.18866947293281555, "diag_loss": 0.1756717711687088, "jvp_rms": 2.2212889194488525, "grad_norm": 0.21811997890472412, "grad_norm_postclip": 0.21811997890472412, "clip_fraction": 0.0, "ema_updates": 48100, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4193356590611594, "updates_nfe_1": 7, "loss_nfe_3": 0.3747627162016355, "updates_nfe_3": 13, "loss_nfe_2": 0.3598556054963006, "updates_nfe_2": 9, "loss_nfe_4": 0.29116837544874713, "updates_nfe_4": 11, "loss_nfe_5": 0.3968243062496185, "updates_nfe_5": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.509881037578676, "grad_norm": 0.04712526251994253, "grad_rms": 0.00013722416146165413, "grad_to_weight": 0.002854367176400556, "update_to_weight": 2.85912292841095e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.258400022421103, "grad_norm": 0.015746224735534258, "grad_rms": 4.5851468405211555e-05, "grad_to_weight": 0.000968497805061966, "update_to_weight": 3.533768166058287e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065746632079983, "grad_norm": 0.0013049155676233698, "grad_rms": 2.3361736183517615e-05, "grad_to_weight": 0.0002575959009397939, "update_to_weight": 1.686898129876448e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866371003807453, "grad_norm": 0.01429020982539393, "grad_rms": 4.57287299740724e-05, "grad_to_weight": 0.0011106635912461354, "update_to_weight": 3.710327074821076e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043279636767022, "grad_norm": 0.015929121589126533, "grad_rms": 0.00036240006447210753, "grad_to_weight": 0.003158484703683384, "update_to_weight": 1.7329901991866458e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644428115879727, "grad_norm": 0.07404727325565792, "grad_rms": 0.0001784276824797142, "grad_to_weight": 0.004448772450464227, "update_to_weight": 4.6427332456217024e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.573336707763662, "grad_norm": 0.01501232222238202, "grad_rms": 2.0711754794696233e-05, "grad_to_weight": 0.000610919160100836, "update_to_weight": 5.128155481057102e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831885501445072, "grad_norm": 0.007046310183277489, "grad_rms": 8.554056409096793e-06, "grad_to_weight": 0.00022853971038998676, "update_to_weight": 4.2035372747925584e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.630571598774004, "grad_norm": 0.003929921414725235, "grad_rms": 3.5278992474753545e-06, "grad_to_weight": 0.00011348127487633969, "update_to_weight": 4.569292864560061e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08207383892266, "grad_norm": 0.003331616829827089, "grad_rms": 2.679825360489854e-06, "grad_to_weight": 8.984440418027822e-05, "update_to_weight": 4.8200234768869004e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67062689805366, "grad_norm": 0.001139529224742965, "grad_rms": 2.0587776138355795e-06, "grad_to_weight": 4.618971497772786e-05, "update_to_weight": 2.9122155714731596e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.354744720383245, "grad_norm": 0.0034045804446492527, "grad_rms": 2.738514506145177e-06, "grad_to_weight": 9.114184744492408e-05, "update_to_weight": 5.2345538967056526e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49556302165924, "grad_norm": 0.009111599463917996, "grad_rms": 4.045234462453923e-06, "grad_to_weight": 0.00014349978219438576, "update_to_weight": 5.800337938927218e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.781271634610775, "grad_norm": 0.020493251734680792, "grad_rms": 1.529147029631993e-05, "grad_to_weight": 0.00040355924684472245, "update_to_weight": 4.817150932917782e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10578206297457, "grad_norm": 0.029057838889090182, "grad_rms": 2.4953725213666863e-05, "grad_to_weight": 0.0007245299154985488, "update_to_weight": 6.196158228925728e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204337967337246, "grad_norm": 0.061248670957096005, "grad_rms": 9.547799855015236e-05, "grad_to_weight": 0.0021716046314586993, "update_to_weight": 5.134504064751999e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.415105643467744, "grad_norm": 0.007601217212098996, "grad_rms": 0.0006470585793515078, "grad_to_weight": 0.0009032824463706488, "update_to_weight": 3.621905456713627e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.321012367886582, "grad_norm": 0.10862847140162504, "grad_rms": 0.002514709651176043, "grad_to_weight": 0.08223122965564207, "update_to_weight": 0.0001371114144723142, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48100}
{"step": 48150, "lr": 1.0336794819816292e-05, "window_steps": 50, "elapsed_seconds": 4132.852723374, "steps_per_second": 1.7300398728369992, "eta_seconds": 1069.3395158380279, "images_seen": 49305600, "images_per_second": 1771.5608297850872, "effective_batch_size": 1024, "loss": 0.38637295365333557, "edge_loss": 0.20120692253112793, "diag_loss": 0.18516594171524048, "jvp_rms": 2.5326952934265137, "grad_norm": 0.24674201011657715, "grad_norm_postclip": 0.24674201011657715, "clip_fraction": 0.0, "ema_updates": 48150, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4097520435849826, "updates_nfe_1": 12, "loss_nfe_3": 0.37426663296563284, "updates_nfe_3": 7, "loss_nfe_5": 0.3882764220237732, "updates_nfe_5": 15, "loss_nfe_4": 0.37283821403980255, "updates_nfe_4": 8, "loss_nfe_2": 0.3718629814684391, "updates_nfe_2": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50784577611412, "grad_norm": 0.0038785913344059294, "grad_rms": 1.1294079121385463e-05, "grad_to_weight": 0.00023495442028044765, "update_to_weight": 2.8192650198732898e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.260771482589146, "grad_norm": 0.0024702802907505495, "grad_rms": 7.193214920130035e-06, "grad_to_weight": 0.00015191654918682957, "update_to_weight": 3.683731333743311e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.0663482733286465, "grad_norm": 0.0006743781089871119, "grad_rms": 1.2073304864306415e-05, "grad_to_weight": 0.00013310930725732325, "update_to_weight": 2.0015861083370274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866596331099416, "grad_norm": 0.00444881146800991, "grad_rms": 1.423621491999847e-05, "grad_to_weight": 0.00034576443944672747, "update_to_weight": 3.420475479540928e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0430735156865545, "grad_norm": 0.01203689444629512, "grad_rms": 0.00027384883083314467, "grad_to_weight": 0.0023868171679144044, "update_to_weight": 1.8381525580698648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64382855020424, "grad_norm": 0.05182294526091383, "grad_rms": 0.00012487492942856487, "grad_to_weight": 0.00311364330055406, "update_to_weight": 4.617380672326636e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.573607440039964, "grad_norm": 0.020069814400702646, "grad_rms": 2.768932537450292e-05, "grad_to_weight": 0.000816722349360929, "update_to_weight": 5.080028813618409e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83205970676819, "grad_norm": 0.00896451868907265, "grad_rms": 1.088271684784985e-05, "grad_to_weight": 0.00029075315675730797, "update_to_weight": 3.831232915999512e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63214896866061, "grad_norm": 0.00609303796993441, "grad_rms": 5.469733819212579e-06, "grad_to_weight": 0.00017593589053477835, "update_to_weight": 4.397687548960111e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08414273753844, "grad_norm": 0.0051461082632997205, "grad_rms": 4.139332983418962e-06, "grad_to_weight": 0.0001387684299383998, "update_to_weight": 4.815254592002107e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67134434791898, "grad_norm": 0.001719153036140484, "grad_rms": 3.105978950528314e-06, "grad_to_weight": 6.968217912638774e-05, "update_to_weight": 2.9720753445975203e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35622013289336, "grad_norm": 0.005167379714380889, "grad_rms": 4.156442926420629e-06, "grad_to_weight": 0.00013832715665552158, "update_to_weight": 5.1523192200006955e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.498091060750596, "grad_norm": 0.013497818310994765, "grad_rms": 5.9925636564470745e-06, "grad_to_weight": 0.00021257045819032548, "update_to_weight": 6.134017039841564e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.781790862504295, "grad_norm": 0.022003451595590105, "grad_rms": 1.6418337648245376e-05, "grad_to_weight": 0.0004332941241707325, "update_to_weight": 5.341712462184027e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.106103172943016, "grad_norm": 0.022919742443746036, "grad_rms": 1.9682570238352766e-05, "grad_to_weight": 0.0005714776712390373, "update_to_weight": 6.592704733154974e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.203924633747228, "grad_norm": 0.044137583106645895, "grad_rms": 6.880423738199222e-05, "grad_to_weight": 0.0015649447259490031, "update_to_weight": 5.587477238938185e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.415335017951184, "grad_norm": 0.005382289110218192, "grad_rms": 0.00045817087555049304, "grad_to_weight": 0.0006395810860455293, "update_to_weight": 5.171969781734058e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3213157679304153, "grad_norm": 0.0695140625841756, "grad_rms": 0.00160922529625384, "grad_to_weight": 0.052609727569554315, "update_to_weight": 0.0001720885969092108, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48150}
{"step": 48200, "lr": 1.0318866508300042e-05, "window_steps": 50, "elapsed_seconds": 4160.429861377, "steps_per_second": 1.730590405294557, "eta_seconds": 1040.10746534425, "images_seen": 49356800, "images_per_second": 1772.1245750216265, "effective_batch_size": 1024, "loss": 0.36599552631378174, "edge_loss": 0.1869119256734848, "diag_loss": 0.17908355593681335, "jvp_rms": 2.0521483421325684, "grad_norm": 0.21240372955799103, "grad_norm_postclip": 0.21240372955799103, "clip_fraction": 0.0, "ema_updates": 48200, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.35935188035170235, "updates_nfe_5": 15, "loss_nfe_1": 0.39173456430435183, "updates_nfe_1": 5, "loss_nfe_3": 0.3163459238906701, "updates_nfe_3": 12, "loss_nfe_4": 0.44325874911414254, "updates_nfe_4": 9, "loss_nfe_2": 0.35170500808291966, "updates_nfe_2": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.510283993289228, "grad_norm": 0.08655123307222537, "grad_rms": 0.0002520287367477735, "grad_to_weight": 0.00524226192035249, "update_to_weight": 4.554212305299681e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.257814377102246, "grad_norm": 0.031022925073300997, "grad_rms": 9.033572762528233e-05, "grad_to_weight": 0.0019081854641540353, "update_to_weight": 5.102979095203433e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066140455392994, "grad_norm": 0.00325668008013532, "grad_rms": 5.830392613432087e-05, "grad_to_weight": 0.000642832568265755, "update_to_weight": 2.301627088448169e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.8664314867838, "grad_norm": 0.03389013462824285, "grad_rms": 0.00010844856962463507, "grad_to_weight": 0.0026339964319597297, "update_to_weight": 5.3941311122336814e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043329369307475, "grad_norm": 0.01500962530019635, "grad_rms": 0.0003414807995568578, "grad_to_weight": 0.002976134255982055, "update_to_weight": 1.7126399092825375e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644409780858105, "grad_norm": 0.097299611232173, "grad_rms": 0.00023445757520864895, "grad_to_weight": 0.005845783209691963, "update_to_weight": 5.2334574002882136e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.573901764864665, "grad_norm": 0.019182045036861814, "grad_rms": 2.646451410908117e-05, "grad_to_weight": 0.0007805860550923161, "update_to_weight": 6.485221515199059e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.831847888802923, "grad_norm": 0.008247129787769276, "grad_rms": 1.0011823434220037e-05, "grad_to_weight": 0.0002674873662296561, "update_to_weight": 5.4938358096020916e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63183350043082, "grad_norm": 0.003749548831279253, "grad_rms": 3.3659783757195113e-06, "grad_to_weight": 0.00010826885129349586, "update_to_weight": 6.291801243523281e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08374443826432, "grad_norm": 0.0031343872760755034, "grad_rms": 2.5211814386409455e-06, "grad_to_weight": 8.452186594300148e-05, "update_to_weight": 6.648408118896948e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671468044062937, "grad_norm": 0.001167514919439704, "grad_rms": 2.1093391268694197e-06, "grad_to_weight": 4.7322474582969154e-05, "update_to_weight": 4.366893599129285e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35636064533137, "grad_norm": 0.0032885959533548066, "grad_rms": 2.6452210102029972e-06, "grad_to_weight": 8.803309253214955e-05, "update_to_weight": 7.067588888534089e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49812566438969, "grad_norm": 0.008512286038941006, "grad_rms": 3.779160066830113e-06, "grad_to_weight": 0.00013405570557990142, "update_to_weight": 8.067506977932614e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78180768924537, "grad_norm": 0.021038319991550036, "grad_rms": 1.569818442677087e-05, "grad_to_weight": 0.00041428852080832004, "update_to_weight": 6.210012112027949e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10636949442616, "grad_norm": 0.042663354559293334, "grad_rms": 3.663760510302632e-05, "grad_to_weight": 0.0010637550867131599, "update_to_weight": 6.699153348416588e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204239503602334, "grad_norm": 0.11576255882992005, "grad_rms": 0.00018045742464954675, "grad_to_weight": 0.004104438228697303, "update_to_weight": 4.816411513383978e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41565549712493, "grad_norm": 0.0032095102572929912, "grad_rms": 0.00027321165670578506, "grad_to_weight": 0.0003813737692078136, "update_to_weight": 2.607494599151966e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3215097720924733, "grad_norm": 0.04320421239491589, "grad_rms": 0.001000161822026056, "grad_to_weight": 0.03269307068876722, "update_to_weight": 5.235290340853965e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48200}
{"step": 48250, "lr": 1.0301426822501102e-05, "window_steps": 50, "elapsed_seconds": 4187.948298593, "steps_per_second": 1.7311579520778084, "eta_seconds": 1010.8840720741724, "images_seen": 49408000, "images_per_second": 1772.7057429276758, "effective_batch_size": 1024, "loss": 0.3547986149787903, "edge_loss": 0.18320675194263458, "diag_loss": 0.17159178853034973, "jvp_rms": 1.9689843654632568, "grad_norm": 0.2364054173231125, "grad_norm_postclip": 0.2343839854001999, "clip_fraction": 0.019999999552965164, "ema_updates": 48250, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3412773609161377, "updates_nfe_4": 11, "loss_nfe_5": 0.3170762062072754, "updates_nfe_5": 13, "loss_nfe_3": 0.3783131897449493, "updates_nfe_3": 10, "loss_nfe_1": 0.39976749155256486, "updates_nfe_1": 9, "loss_nfe_2": 0.3546924037592752, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50748250823746, "grad_norm": 0.09330298724011171, "grad_rms": 0.0002716891854018548, "grad_to_weight": 0.005652163326149354, "update_to_weight": 3.697560321919578e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2598236907747, "grad_norm": 0.15257905247841486, "grad_rms": 0.0004442952975403336, "grad_to_weight": 0.00938380731428123, "update_to_weight": 6.189875862473265e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.065853938789628, "grad_norm": 0.004929184877597716, "grad_rms": 8.824656519344915e-05, "grad_to_weight": 0.0009730215156529828, "update_to_weight": 2.5696734333527077e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866502049896857, "grad_norm": 0.061061193634503776, "grad_rms": 0.0001953960697375414, "grad_to_weight": 0.004745749341795136, "update_to_weight": 5.360201888576204e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0429129623294, "grad_norm": 0.028736068672883985, "grad_rms": 0.0006537681994239276, "grad_to_weight": 0.005698307483698936, "update_to_weight": 1.8671520293301442e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644362109707373, "grad_norm": 0.16563490801991154, "grad_rms": 0.00039912141901154, "grad_to_weight": 0.009951412191597868, "update_to_weight": 6.219102224928393e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.573898039256, "grad_norm": 0.030239831630104564, "grad_rms": 4.1720392653299056e-05, "grad_to_weight": 0.0012305671481910366, "update_to_weight": 8.995789919270312e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832069604768364, "grad_norm": 0.008916876685783883, "grad_rms": 1.0824880565742705e-05, "grad_to_weight": 0.00028920785403276447, "update_to_weight": 8.936608022058384e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.632156018197165, "grad_norm": 0.00299830791571945, "grad_rms": 2.6915877248669104e-06, "grad_to_weight": 8.657583761588551e-05, "update_to_weight": 0.00010455450818257762, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.083718104277246, "grad_norm": 0.0024245446349274273, "grad_rms": 1.950211123364792e-06, "grad_to_weight": 6.538030054348245e-05, "update_to_weight": 0.00010496070337243648, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.670960885931382, "grad_norm": 0.0010056176300367878, "grad_rms": 1.8168406916155401e-06, "grad_to_weight": 4.0761186185100775e-05, "update_to_weight": 7.04058749808069e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35652076351179, "grad_norm": 0.0025260618464900087, "grad_rms": 2.031867691922149e-06, "grad_to_weight": 6.762037242390505e-05, "update_to_weight": 0.00010188976958444288, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49878697475744, "grad_norm": 0.007514827020421216, "grad_rms": 3.336322822657981e-06, "grad_to_weight": 0.00011834599334013376, "update_to_weight": 0.00010146957244043608, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78205768588449, "grad_norm": 0.02797351812497968, "grad_rms": 2.0873028206050913e-05, "grad_to_weight": 0.0005508543647051795, "update_to_weight": 6.410349646968044e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10661603045141, "grad_norm": 0.06563157988100198, "grad_rms": 5.636181052350204e-05, "grad_to_weight": 0.0016364277612244931, "update_to_weight": 6.374992632001883e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20431199891306, "grad_norm": 0.240273131368954, "grad_rms": 0.0003745517630016058, "grad_to_weight": 0.008519021182938753, "update_to_weight": 5.314740979628985e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41543157095289, "grad_norm": 0.006776301029096493, "grad_rms": 0.0005768370505405618, "grad_to_weight": 0.0008052232344786571, "update_to_weight": 2.7362755845482212e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3213585315833236, "grad_norm": 0.13365341884030682, "grad_rms": 0.0030940280934981987, "grad_to_weight": 0.1011484889571614, "update_to_weight": 8.055677323613126e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48250}
{"step": 48300, "lr": 1.028447595313709e-05, "window_steps": 50, "elapsed_seconds": 4215.458525667, "steps_per_second": 1.7317214617465468, "eta_seconds": 981.6821224156026, "images_seen": 49459200, "images_per_second": 1773.282776828464, "effective_batch_size": 1024, "loss": 0.3729509711265564, "edge_loss": 0.1925518959760666, "diag_loss": 0.1803990751504898, "jvp_rms": 2.3388280868530273, "grad_norm": 0.19813866913318634, "grad_norm_postclip": 0.19813866913318634, "clip_fraction": 0.0, "ema_updates": 48300, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.36426417032877606, "updates_nfe_4": 9, "loss_nfe_1": 0.3931449979543686, "updates_nfe_1": 10, "loss_nfe_2": 0.35413446128368375, "updates_nfe_2": 10, "loss_nfe_5": 0.37854187852806515, "updates_nfe_5": 9, "loss_nfe_3": 0.3741248995065689, "updates_nfe_3": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50910282463038, "grad_norm": 0.018871776589628395, "grad_rms": 5.4952770113642004e-05, "grad_to_weight": 0.0011431133956881702, "update_to_weight": 4.568075868035943e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25887021277011, "grad_norm": 0.014813392427678188, "grad_rms": 4.3135151839850913e-05, "grad_to_weight": 0.0009110960499606788, "update_to_weight": 7.62670468033189e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06630591982144, "grad_norm": 0.002579546587778757, "grad_rms": 4.6181291994650834e-05, "grad_to_weight": 0.0005091572890785229, "update_to_weight": 2.7512200222232735e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866066805703603, "grad_norm": 0.014225814235786914, "grad_rms": 4.552266382356511e-05, "grad_to_weight": 0.001105684779242754, "update_to_weight": 6.841988896334743e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.043019431145408, "grad_norm": 0.04066369089511958, "grad_rms": 0.0009251310010794674, "grad_to_weight": 0.008063361930351265, "update_to_weight": 2.312480347769445e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64396423362404, "grad_norm": 0.061736898379600794, "grad_rms": 0.000148764042442523, "grad_to_weight": 0.0037092664651898416, "update_to_weight": 5.261563017259068e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.573337949661646, "grad_norm": 0.02589447637223583, "grad_rms": 3.572532198644117e-05, "grad_to_weight": 0.0010537630836022574, "update_to_weight": 7.025816432590002e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83225766616789, "grad_norm": 0.015913440237171558, "grad_rms": 1.9318545722640875e-05, "grad_to_weight": 0.0005161295812156276, "update_to_weight": 6.060342312902554e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.633313885102915, "grad_norm": 0.01112382729064127, "grad_rms": 9.98588464915727e-06, "grad_to_weight": 0.000321188648812092, "update_to_weight": 7.118360946282192e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08449824071542, "grad_norm": 0.0091568846079613, "grad_rms": 7.365448324010132e-06, "grad_to_weight": 0.0002469194688444745, "update_to_weight": 7.666214542554114e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671790888077016, "grad_norm": 0.003545168539337548, "grad_rms": 6.405025397842377e-06, "grad_to_weight": 0.00014369319825302182, "update_to_weight": 5.6992304858186566e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35721841328828, "grad_norm": 0.008781830391733522, "grad_rms": 7.063769033880598e-06, "grad_to_weight": 0.00023507720233821664, "update_to_weight": 7.437253037810535e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49995963027752, "grad_norm": 0.020852375377481906, "grad_rms": 9.257732172606245e-06, "grad_to_weight": 0.00032838407298040626, "update_to_weight": 8.438586493225627e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78245431264252, "grad_norm": 0.02850124812392993, "grad_rms": 2.1266805031120026e-05, "grad_to_weight": 0.0005612420374261906, "update_to_weight": 6.743987261696946e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10657798486739, "grad_norm": 0.03084492971544514, "grad_rms": 2.6488408284315642e-05, "grad_to_weight": 0.000769074083734674, "update_to_weight": 7.211818736730704e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20410966078902, "grad_norm": 0.06821119051038642, "grad_rms": 0.00010633157988386214, "grad_to_weight": 0.0024184840908209044, "update_to_weight": 6.425315989823289e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.415986842533169, "grad_norm": 0.008308907170867578, "grad_rms": 0.0007073011492669223, "grad_to_weight": 0.0009872766350911546, "update_to_weight": 3.982018714606224e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.321939132252962, "grad_norm": 0.09327852387334534, "grad_rms": 0.0021593639421151555, "grad_to_weight": 0.07056189017899189, "update_to_weight": 9.157673636742829e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48300}
{"step": 48350, "lr": 1.0268014085579995e-05, "window_steps": 50, "elapsed_seconds": 4242.988391448, "steps_per_second": 1.7322696462743972, "eta_seconds": 952.5075980801632, "images_seen": 49510400, "images_per_second": 1773.8441177849827, "effective_batch_size": 1024, "loss": 0.3922787010669708, "edge_loss": 0.2030271291732788, "diag_loss": 0.1892516165971756, "jvp_rms": 2.523984432220459, "grad_norm": 0.3157505989074707, "grad_norm_postclip": 0.3141879439353943, "clip_fraction": 0.019999999552965164, "ema_updates": 48350, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3564737066626549, "updates_nfe_3": 6, "loss_nfe_5": 0.4012596729923697, "updates_nfe_5": 17, "loss_nfe_4": 0.3942488594488664, "updates_nfe_4": 11, "loss_nfe_1": 0.42240012543542044, "updates_nfe_1": 7, "loss_nfe_2": 0.3733490738603804, "updates_nfe_2": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.509048292894093, "grad_norm": 0.044191564593793926, "grad_rms": 0.00012868151965192046, "grad_to_weight": 0.002676808729962652, "update_to_weight": 3.307445745327241e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26125474130821, "grad_norm": 0.012792845922892367, "grad_rms": 3.725151777635523e-05, "grad_to_weight": 0.0007867071838186571, "update_to_weight": 5.595745114067616e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066287660644665, "grad_norm": 0.004225578834829892, "grad_rms": 7.564999637619754e-05, "grad_to_weight": 0.0008340582134043694, "update_to_weight": 2.643371425784245e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.867757889983569, "grad_norm": 0.016525020227653422, "grad_rms": 5.288013241510376e-05, "grad_to_weight": 0.0012842190822160802, "update_to_weight": 5.341715410641203e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.04271855165796, "grad_norm": 0.06391424259672868, "grad_rms": 0.0014540993681378823, "grad_to_weight": 0.012674560743770791, "update_to_weight": 1.6363049659761044e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644928654168023, "grad_norm": 0.10884632336347853, "grad_rms": 0.00026228105871135615, "grad_to_weight": 0.0065393084959978216, "update_to_weight": 5.497057329933597e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574020984043674, "grad_norm": 0.03721228716365035, "grad_rms": 5.133994299258089e-05, "grad_to_weight": 0.0015142937815432367, "update_to_weight": 6.94550924816676e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83201912493427, "grad_norm": 0.01733423202633243, "grad_rms": 2.1043353855400454e-05, "grad_to_weight": 0.0005622152722496855, "update_to_weight": 5.173499015217954e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63190928356391, "grad_norm": 0.01176159323655758, "grad_rms": 1.0558408565852653e-05, "grad_to_weight": 0.0003396172339288137, "update_to_weight": 5.754921280223049e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.083676957385, "grad_norm": 0.009168335815024952, "grad_rms": 7.374659237709049e-06, "grad_to_weight": 0.000247233731044546, "update_to_weight": 5.862040827201805e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671329504340026, "grad_norm": 0.003684582647418095, "grad_rms": 6.656903663477754e-06, "grad_to_weight": 0.00014934674058687944, "update_to_weight": 3.706591838623948e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35684099781374, "grad_norm": 0.009305669182455636, "grad_rms": 7.4851249544107265e-06, "grad_to_weight": 0.0002491021439152266, "update_to_weight": 6.525178117366342e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.49906764509913, "grad_norm": 0.022391616577718376, "grad_rms": 9.94110193374235e-06, "grad_to_weight": 0.00035262906067954784, "update_to_weight": 7.927046676285149e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78212499246167, "grad_norm": 0.03172561319331062, "grad_rms": 2.3672732763881287e-05, "grad_to_weight": 0.000624739772075708, "update_to_weight": 6.360670783570964e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.107179100875285, "grad_norm": 0.035821910966008015, "grad_rms": 3.076244335602729e-05, "grad_to_weight": 0.0008931545865120753, "update_to_weight": 7.535681283893824e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204831362853977, "grad_norm": 0.07615733924970382, "grad_rms": 0.00011871849972973586, "grad_to_weight": 0.002700152263629689, "update_to_weight": 6.29714741511889e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.415441996797604, "grad_norm": 0.009266991863057367, "grad_rms": 0.0007888587343915767, "grad_to_weight": 0.0011011889650696673, "update_to_weight": 3.869757309110757e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.321435123780844, "grad_norm": 0.10665398887765706, "grad_rms": 0.002469001098022027, "grad_to_weight": 0.08071072651111498, "update_to_weight": 7.385522183718316e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48350}
{"step": 48400, "lr": 1.0252041399854163e-05, "window_steps": 50, "elapsed_seconds": 4270.490367908999, "steps_per_second": 1.7328220795457108, "eta_seconds": 923.3492687370809, "images_seen": 49561600, "images_per_second": 1774.4098094548078, "effective_batch_size": 1024, "loss": 0.36110106110572815, "edge_loss": 0.1879347413778305, "diag_loss": 0.1731664091348648, "jvp_rms": 1.9697265625, "grad_norm": 0.22041749954223633, "grad_norm_postclip": 0.22041749954223633, "clip_fraction": 0.0, "ema_updates": 48400, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.31636928021907806, "updates_nfe_4": 6, "loss_nfe_1": 0.40041569620370865, "updates_nfe_1": 12, "loss_nfe_2": 0.34618550539016724, "updates_nfe_2": 6, "loss_nfe_5": 0.3163949708853449, "updates_nfe_5": 14, "loss_nfe_3": 0.4037676180402438, "updates_nfe_3": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.507943755393946, "grad_norm": 0.056556795831074695, "grad_rms": 0.00016468786523136944, "grad_to_weight": 0.0034260351664085874, "update_to_weight": 6.07568187763254e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262522406126475, "grad_norm": 0.038509584854223296, "grad_rms": 0.00011213614964205089, "grad_to_weight": 0.0023679958060943748, "update_to_weight": 6.864904922086915e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067405113984218, "grad_norm": 0.0032929031848478586, "grad_rms": 5.895242373603429e-05, "grad_to_weight": 0.000649820393431863, "update_to_weight": 2.893515054820055e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.867722315464016, "grad_norm": 0.023178954419499975, "grad_rms": 7.41727490835795e-05, "grad_to_weight": 0.0018013253512351793, "update_to_weight": 6.136502056998629e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.042848096566942, "grad_norm": 0.019934806413447886, "grad_rms": 0.0004535325497423516, "grad_to_weight": 0.003953084850408057, "update_to_weight": 2.217744109697316e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64420717752454, "grad_norm": 0.03648376305299737, "grad_rms": 8.791293728277752e-05, "grad_to_weight": 0.0021919796277387804, "update_to_weight": 5.4609424918391114e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.573386383634123, "grad_norm": 0.03690759959234371, "grad_rms": 5.091958069470228e-05, "grad_to_weight": 0.001501933800093754, "update_to_weight": 6.98790471348622e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83239029825463, "grad_norm": 0.012646532679984729, "grad_rms": 1.535259605528144e-05, "grad_to_weight": 0.0004101703616764551, "update_to_weight": 5.683408913181412e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.634406509391965, "grad_norm": 0.009005281914495556, "grad_rms": 8.084061724596472e-06, "grad_to_weight": 0.00026000970774693523, "update_to_weight": 6.771867070007303e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08536394253129, "grad_norm": 0.007282408197672407, "grad_rms": 5.857691076249781e-06, "grad_to_weight": 0.00019636879414093033, "update_to_weight": 7.135167165262126e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67159173958672, "grad_norm": 0.0029525401054131322, "grad_rms": 5.334328721887241e-06, "grad_to_weight": 0.00011967367718215131, "update_to_weight": 4.6919364623159725e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35824934261926, "grad_norm": 0.00717320703879554, "grad_rms": 5.769853833883498e-06, "grad_to_weight": 0.0001920113272174175, "update_to_weight": 7.544806094430414e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.501326419463474, "grad_norm": 0.017979287797113056, "grad_rms": 7.98218083392187e-06, "grad_to_weight": 0.0002831324762942639, "update_to_weight": 8.25623156012321e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.782730747641956, "grad_norm": 0.026050398649113703, "grad_rms": 1.943805220896633e-05, "grad_to_weight": 0.0005129775076210002, "update_to_weight": 6.751404110798838e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10738149993402, "grad_norm": 0.03220879903259673, "grad_rms": 2.7659645426122847e-05, "grad_to_weight": 0.0008030641200710076, "update_to_weight": 7.350834396913257e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204251405830313, "grad_norm": 0.05810499191088522, "grad_rms": 9.057744840390493e-05, "grad_to_weight": 0.0020601501197395336, "update_to_weight": 5.854363737902947e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416005879766347, "grad_norm": 0.007437988642235096, "grad_rms": 0.0006331636407412023, "grad_to_weight": 0.0008837908086682083, "update_to_weight": 6.6467279567534905e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.32209963862506, "grad_norm": 0.09584786699667305, "grad_rms": 0.0022188433020476593, "grad_to_weight": 0.07249670463290624, "update_to_weight": 0.0001086858166982492, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48400}
{"step": 48450, "lr": 1.0236558070634328e-05, "window_steps": 50, "elapsed_seconds": 4298.077227379999, "steps_per_second": 1.7333332106136525, "eta_seconds": 894.2308325421475, "images_seen": 49612800, "images_per_second": 1774.9332076683802, "effective_batch_size": 1024, "loss": 0.3886876702308655, "edge_loss": 0.2013210654258728, "diag_loss": 0.18736675381660461, "jvp_rms": 2.497617244720459, "grad_norm": 0.17583289742469788, "grad_norm_postclip": 0.17583289742469788, "clip_fraction": 0.0, "ema_updates": 48450, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4020906388759613, "updates_nfe_1": 10, "loss_nfe_3": 0.3998602216060345, "updates_nfe_3": 13, "loss_nfe_5": 0.3768211986337389, "updates_nfe_5": 7, "loss_nfe_4": 0.4248269349336624, "updates_nfe_4": 8, "loss_nfe_2": 0.34824445346991223, "updates_nfe_2": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.508958638631064, "grad_norm": 0.08987483751410968, "grad_rms": 0.0002617067482469094, "grad_to_weight": 0.005444004039346371, "update_to_weight": 3.5782548326651564e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.260699227138687, "grad_norm": 0.0664627837311082, "grad_rms": 0.000193533134421247, "grad_to_weight": 0.004087326307603275, "update_to_weight": 5.47254954644701e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06673828542854, "grad_norm": 0.003269137133829669, "grad_rms": 5.852694317025165e-05, "grad_to_weight": 0.0006452153140081061, "update_to_weight": 2.6594551829894063e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.8678397110053, "grad_norm": 0.030701686546976484, "grad_rms": 9.824552270467428e-05, "grad_to_weight": 0.0023859239185827497, "update_to_weight": 4.949866522001843e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.042511083754319, "grad_norm": 0.030158997547044455, "grad_rms": 0.0006861409522370507, "grad_to_weight": 0.005980948191508994, "update_to_weight": 1.7183569445251256e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644880984503352, "grad_norm": 0.11369458378091891, "grad_rms": 0.0002739636478508005, "grad_to_weight": 0.006830603588380738, "update_to_weight": 4.775255460390394e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57436870334554, "grad_norm": 0.015614855069395982, "grad_rms": 2.1543039415305933e-05, "grad_to_weight": 0.00063541225648129, "update_to_weight": 5.894766686787931e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832353676019043, "grad_norm": 0.007504571239330181, "grad_rms": 9.110374655328706e-06, "grad_to_weight": 0.00024339923309737874, "update_to_weight": 4.751019207727048e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63329273719599, "grad_norm": 0.0036906618167976105, "grad_rms": 3.31311537105549e-06, "grad_to_weight": 0.00010656398872619634, "update_to_weight": 5.418869283345167e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08487678201857, "grad_norm": 0.002959228660601825, "grad_rms": 2.3802905367664296e-06, "grad_to_weight": 7.979610335490377e-05, "update_to_weight": 5.787884139358314e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67103386786567, "grad_norm": 0.0011077100635012987, "grad_rms": 2.0012902098858094e-06, "grad_to_weight": 4.489921538894667e-05, "update_to_weight": 3.734009522123802e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35801080999479, "grad_norm": 0.0031727542000202173, "grad_rms": 2.5520423272252928e-06, "grad_to_weight": 8.492834953544627e-05, "update_to_weight": 6.28933765491699e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.501139953473476, "grad_norm": 0.008346064638437031, "grad_rms": 3.7053635242605274e-06, "grad_to_weight": 0.00013143172932882926, "update_to_weight": 7.573985707523262e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78262017382276, "grad_norm": 0.01984715758304229, "grad_rms": 1.4809373572173089e-05, "grad_to_weight": 0.0003908257887266918, "update_to_weight": 6.035102119347336e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.107465198494374, "grad_norm": 0.04586963561523732, "grad_rms": 3.9391032731735877e-05, "grad_to_weight": 0.0011436682769211566, "update_to_weight": 6.306141176798573e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204658242602964, "grad_norm": 0.10450862594787529, "grad_rms": 0.00016291413806708262, "grad_to_weight": 0.003705367569035658, "update_to_weight": 5.090630194731636e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.415753406291548, "grad_norm": 0.007563176322288613, "grad_rms": 0.0006438203237628116, "grad_to_weight": 0.0008986927203254844, "update_to_weight": 3.583551571635417e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3216908049253266, "grad_norm": 0.08185167753031837, "grad_rms": 0.0018948366003368217, "grad_to_weight": 0.06192952029725504, "update_to_weight": 6.370541458357168e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48450}
{"step": 48500, "lr": 1.0221564267243686e-05, "window_steps": 50, "elapsed_seconds": 4325.569566183, "steps_per_second": 1.733875709371195, "eta_seconds": 865.1139132365998, "images_seen": 49664000, "images_per_second": 1775.4887263961036, "effective_batch_size": 1024, "loss": 0.3684673607349396, "edge_loss": 0.18998336791992188, "diag_loss": 0.1784840226173401, "jvp_rms": 2.2960546016693115, "grad_norm": 0.21510644257068634, "grad_norm_postclip": 0.21510644257068634, "clip_fraction": 0.0, "ema_updates": 48500, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39556122819582623, "updates_nfe_1": 9, "loss_nfe_4": 0.3717641197144985, "updates_nfe_4": 8, "loss_nfe_2": 0.3652500092983246, "updates_nfe_2": 15, "loss_nfe_3": 0.34473155438899994, "updates_nfe_3": 11, "loss_nfe_5": 0.3740582082952772, "updates_nfe_5": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.508883772327042, "grad_norm": 0.042906691971764804, "grad_rms": 0.00012494009607749753, "grad_to_weight": 0.002599006241941505, "update_to_weight": 2.8795947496642954e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.259132987880324, "grad_norm": 0.0467781989630645, "grad_rms": 0.00013621354628372634, "grad_to_weight": 0.0028770414140737587, "update_to_weight": 3.765521682406711e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066644173287773, "grad_norm": 0.0031988737239456405, "grad_rms": 5.72690263472846e-05, "grad_to_weight": 0.0006313594589512834, "update_to_weight": 2.1437211428912672e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.866859010303004, "grad_norm": 0.018869339775005458, "grad_rms": 6.0381964568981576e-05, "grad_to_weight": 0.0014665070752618047, "update_to_weight": 4.198713394297421e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.04274389355162, "grad_norm": 0.032557246672893614, "grad_rms": 0.0007407030090940459, "grad_to_weight": 0.006456256228781717, "update_to_weight": 1.6849690999140185e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644536292096163, "grad_norm": 0.07498816034126739, "grad_rms": 0.00018069488685847318, "grad_to_weight": 0.00450527182165335, "update_to_weight": 4.8655591757146405e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574383605491384, "grad_norm": 0.03479398236597561, "grad_rms": 4.800352806314373e-05, "grad_to_weight": 0.0014158638899980612, "update_to_weight": 6.565370920616111e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832538766331645, "grad_norm": 0.02460290354456176, "grad_rms": 2.9867351744918335e-05, "grad_to_weight": 0.0007979525698813848, "update_to_weight": 6.087501439133009e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63348306789345, "grad_norm": 0.016866371529780773, "grad_rms": 1.5140979461980803e-05, "grad_to_weight": 0.0004869961099990124, "update_to_weight": 7.265730753754154e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.085337609694214, "grad_norm": 0.014680148113883811, "grad_rms": 1.1808150582963314e-05, "grad_to_weight": 0.0003958477678802735, "update_to_weight": 7.740011376721647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671693169453455, "grad_norm": 0.005129173626743363, "grad_rms": 9.26683371599968e-06, "grad_to_weight": 0.00020789710667664678, "update_to_weight": 4.982294884455933e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35840945270472, "grad_norm": 0.013231025469039466, "grad_rms": 1.06425316620397e-05, "grad_to_weight": 0.0003541645820277702, "update_to_weight": 8.038067138099085e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.502202994512324, "grad_norm": 0.025432735149155533, "grad_rms": 1.1291253210508014e-05, "grad_to_weight": 0.00040050161962654045, "update_to_weight": 8.523448476308647e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.783062467655, "grad_norm": 0.031299676377797944, "grad_rms": 2.3354911061066853e-05, "grad_to_weight": 0.0006163408596662222, "update_to_weight": 5.9133806649779515e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10773455449397, "grad_norm": 0.034638548510137916, "grad_rms": 2.9746218382633303e-05, "grad_to_weight": 0.0008636376223911343, "update_to_weight": 6.102793279565646e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204495941402566, "grad_norm": 0.0757636016519283, "grad_rms": 0.00011810471861086253, "grad_to_weight": 0.002686224274645225, "update_to_weight": 4.2841366326369465e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41621347346638, "grad_norm": 0.0036590109820969345, "grad_rms": 0.0003114756994627998, "grad_to_weight": 0.00043475738746796545, "update_to_weight": 2.1704684416305885e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.322219374508467, "grad_norm": 0.04402036879996292, "grad_rms": 0.0010190555463154526, "grad_to_weight": 0.03329278760291001, "update_to_weight": 4.363454334215448e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48500}
{"step": 48550, "lr": 1.0207060153652086e-05, "window_steps": 50, "elapsed_seconds": 4355.216616948, "steps_per_second": 1.7335532681933064, "eta_seconds": 836.4323304072317, "images_seen": 49715200, "images_per_second": 1775.1585466299457, "effective_batch_size": 1024, "loss": 0.3415820896625519, "edge_loss": 0.17600850760936737, "diag_loss": 0.16557365655899048, "jvp_rms": 1.865078091621399, "grad_norm": 0.16836638748645782, "grad_norm_postclip": 0.16836638748645782, "clip_fraction": 0.0, "ema_updates": 48550, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.330889282318262, "updates_nfe_3": 13, "loss_nfe_1": 0.3983028642833233, "updates_nfe_1": 8, "loss_nfe_2": 0.3735620751976967, "updates_nfe_2": 12, "loss_nfe_4": 0.3086720195909341, "updates_nfe_4": 12, "loss_nfe_5": 0.2808629125356674, "updates_nfe_5": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.508681354323663, "grad_norm": 0.0032669782352329887, "grad_rms": 9.513121516375762e-06, "grad_to_weight": 0.0001978945601477346, "update_to_weight": 2.4728026373039716e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26176801202603, "grad_norm": 0.006926063777246126, "grad_rms": 2.016802121071124e-05, "grad_to_weight": 0.00042591087095352164, "update_to_weight": 3.132926307101047e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066394767659983, "grad_norm": 0.0008729735912471066, "grad_rms": 1.5628734333391827e-05, "grad_to_weight": 0.00017230666603785142, "update_to_weight": 1.6509536008074362e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.867829631632171, "grad_norm": 0.005180568683797481, "grad_rms": 1.6577841007802008e-05, "grad_to_weight": 0.00040259848258034265, "update_to_weight": 3.61385517931206e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.042321196805873, "grad_norm": 0.013681357365139448, "grad_rms": 0.00031126165767841034, "grad_to_weight": 0.0027133054066064036, "update_to_weight": 1.695048140609528e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644432699631974, "grad_norm": 0.028581285204254153, "grad_rms": 6.887076670169146e-05, "grad_to_weight": 0.0017171678794967952, "update_to_weight": 4.6736667301063725e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57429543433036, "grad_norm": 0.01394761778289138, "grad_rms": 1.9242834999817512e-05, "grad_to_weight": 0.000567569386482044, "update_to_weight": 6.222411853151649e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832301217065258, "grad_norm": 0.011090154259230652, "grad_rms": 1.346318891044314e-05, "grad_to_weight": 0.00035969271904662123, "update_to_weight": 5.8731564035055285e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63389016059631, "grad_norm": 0.007914458057310619, "grad_rms": 7.10482670720619e-06, "grad_to_weight": 0.00022851773279327026, "update_to_weight": 7.244289042621672e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08556966717732, "grad_norm": 0.006653067984902643, "grad_rms": 5.351473840384748e-06, "grad_to_weight": 0.00017939775617875864, "update_to_weight": 7.485941594634888e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671810679143515, "grad_norm": 0.0023419439430905965, "grad_rms": 4.231169906134093e-06, "grad_to_weight": 9.492387784372773e-05, "update_to_weight": 5.1576293194646454e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35847316979275, "grad_norm": 0.006827445748287691, "grad_rms": 5.491736654655921e-06, "grad_to_weight": 0.00018275494604014532, "update_to_weight": 7.412146536279745e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50301612239079, "grad_norm": 0.016380587230339763, "grad_rms": 7.272413174196947e-06, "grad_to_weight": 0.0002579497515325743, "update_to_weight": 7.85953668736563e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78308169773424, "grad_norm": 0.02181766670188, "grad_rms": 1.6279710346904742e-05, "grad_to_weight": 0.00042962470910569867, "update_to_weight": 5.665582392011835e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.107947603155914, "grad_norm": 0.02094119749341403, "grad_rms": 1.7983473922142897e-05, "grad_to_weight": 0.0005221208948564164, "update_to_weight": 6.0711588560346065e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204526237697536, "grad_norm": 0.03862946664334454, "grad_rms": 6.02178643639467e-05, "grad_to_weight": 0.0013696194120684524, "update_to_weight": 4.283076763655631e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41609698592896, "grad_norm": 0.004889852507629097, "grad_rms": 0.0004162518881566312, "grad_to_weight": 0.0005810119008614729, "update_to_weight": 2.721949215132915e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.322114110292474, "grad_norm": 0.05467782403291661, "grad_rms": 0.0012657717633944709, "grad_to_weight": 0.04135635767537565, "update_to_weight": 8.061428549841443e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48550}
{"step": 48600, "lr": 1.0193045888474177e-05, "window_steps": 50, "elapsed_seconds": 4382.690037314, "steps_per_second": 1.7340947991516595, "eta_seconds": 807.337638452579, "images_seen": 49766400, "images_per_second": 1775.7130743312994, "effective_batch_size": 1024, "loss": 0.35971227288246155, "edge_loss": 0.18737159669399261, "diag_loss": 0.17234069108963013, "jvp_rms": 2.1206250190734863, "grad_norm": 0.1696002185344696, "grad_norm_postclip": 0.1696002185344696, "clip_fraction": 0.0, "ema_updates": 48600, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.3968943196993608, "updates_nfe_1": 13, "loss_nfe_2": 0.38654763251543045, "updates_nfe_2": 12, "loss_nfe_4": 0.31662722338329663, "updates_nfe_4": 11, "loss_nfe_3": 0.3580630086362362, "updates_nfe_3": 8, "loss_nfe_5": 0.30666891237099964, "updates_nfe_5": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.510704498509728, "grad_norm": 0.026186879548769235, "grad_rms": 7.625363542232805e-05, "grad_to_weight": 0.0015860546441935827, "update_to_weight": 3.49750949056679e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25973172387607, "grad_norm": 0.007623097443623647, "grad_rms": 2.2197715163900798e-05, "grad_to_weight": 0.0004688329163776891, "update_to_weight": 3.060050211059438e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066432602813113, "grad_norm": 0.0015285019382933781, "grad_rms": 2.7364574325249787e-05, "grad_to_weight": 0.0003016919513435715, "update_to_weight": 2.331955274199306e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.868000979901657, "grad_norm": 0.0077858127722800095, "grad_rms": 2.4914632762046375e-05, "grad_to_weight": 0.0006050522365082624, "update_to_weight": 5.2642924010519854e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.042366021361777, "grad_norm": 0.022113868156566435, "grad_rms": 0.0005031079209752463, "grad_to_weight": 0.004385613432837271, "update_to_weight": 2.3912178690305114e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644591296682346, "grad_norm": 0.03962362463157529, "grad_rms": 9.547889076277209e-05, "grad_to_weight": 0.002380570596496004, "update_to_weight": 8.110595998768994e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574312820218136, "grad_norm": 0.018234086135675998, "grad_rms": 2.5156662330657582e-05, "grad_to_weight": 0.0007419978035224727, "update_to_weight": 8.549284237461825e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832279441624266, "grad_norm": 0.013857891600365116, "grad_rms": 1.6823157564365714e-05, "grad_to_weight": 0.00044946049566665035, "update_to_weight": 7.235569772127178e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63444351703048, "grad_norm": 0.010854246651660475, "grad_rms": 9.743881506338506e-06, "grad_to_weight": 0.00031339457342004684, "update_to_weight": 8.173699522126891e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08544129513208, "grad_norm": 0.009202335514558313, "grad_rms": 7.402007297738889e-06, "grad_to_weight": 0.000248138762629912, "update_to_weight": 8.497799951283773e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671846550411082, "grad_norm": 0.0032719041160078855, "grad_rms": 5.9113209230530635e-06, "grad_to_weight": 0.0001326169125331792, "update_to_weight": 5.644260229658193e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35882769622111, "grad_norm": 0.00915144045383905, "grad_rms": 7.361069254318623e-06, "grad_to_weight": 0.0002449605894556678, "update_to_weight": 8.698936812100376e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50356396863683, "grad_norm": 0.020883443028818838, "grad_rms": 9.271525133365037e-06, "grad_to_weight": 0.0003288546614349513, "update_to_weight": 9.202439673862754e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78318505928537, "grad_norm": 0.025656130379950882, "grad_rms": 1.914386065729169e-05, "grad_to_weight": 0.0005052091622453253, "update_to_weight": 7.42264204261139e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10824282585968, "grad_norm": 0.02735160284868842, "grad_rms": 2.3488477042102868e-05, "grad_to_weight": 0.0006819446807341445, "update_to_weight": 8.377697657350775e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20460847033413, "grad_norm": 0.06825965351587908, "grad_rms": 0.00010640712684179378, "grad_to_weight": 0.0024201595844762468, "update_to_weight": 6.241126753219808e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416281461556327, "grad_norm": 0.010403578173511273, "grad_rms": 0.0008856113863460645, "grad_to_weight": 0.001236125267558181, "update_to_weight": 4.937317505297927e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3223150291961399, "grad_norm": 0.12814513112296827, "grad_rms": 0.002966513233853046, "grad_to_weight": 0.09690968361818447, "update_to_weight": 0.0001837411202846678, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48600}
{"step": 48650, "lr": 1.017952162496773e-05, "window_steps": 50, "elapsed_seconds": 4410.151461066999, "steps_per_second": 1.7346343016865788, "eta_seconds": 778.2620225412352, "images_seen": 49817600, "images_per_second": 1776.2655249270567, "effective_batch_size": 1024, "loss": 0.3574724793434143, "edge_loss": 0.18289171159267426, "diag_loss": 0.17458079755306244, "jvp_rms": 2.0721874237060547, "grad_norm": 0.17209281027317047, "grad_norm_postclip": 0.17209281027317047, "clip_fraction": 0.0, "ema_updates": 48650, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3803846801028532, "updates_nfe_2": 17, "loss_nfe_5": 0.3485401049256325, "updates_nfe_5": 10, "loss_nfe_4": 0.3249135845237308, "updates_nfe_4": 9, "loss_nfe_3": 0.3333260764678319, "updates_nfe_3": 9, "loss_nfe_1": 0.3995056629180908, "updates_nfe_1": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.50955755684132, "grad_norm": 0.03924300984628451, "grad_rms": 0.00011427181158108052, "grad_to_weight": 0.0023769873729912756, "update_to_weight": 2.577287938734757e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.261230344124872, "grad_norm": 0.024305654881644333, "grad_rms": 7.077569294173284e-05, "grad_to_weight": 0.001494699624030963, "update_to_weight": 4.624893386715679e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066268271962235, "grad_norm": 0.0029747618250497215, "grad_rms": 5.325677974106987e-05, "grad_to_weight": 0.0005871702139250426, "update_to_weight": 2.1073377373179102e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.868244658113037, "grad_norm": 0.01357712118590152, "grad_rms": 4.344684340688001e-05, "grad_to_weight": 0.0010550872746533893, "update_to_weight": 3.71600106784129e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.042110118857342, "grad_norm": 0.012641313915361633, "grad_rms": 0.000287599849891686, "grad_to_weight": 0.002507147527001343, "update_to_weight": 1.6648370854133807e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644408864106495, "grad_norm": 0.05825018492362851, "grad_rms": 0.00014036229887970473, "grad_to_weight": 0.00349968481303319, "update_to_weight": 4.176257689505036e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574196086164243, "grad_norm": 0.015427079557530218, "grad_rms": 2.128397487481706e-05, "grad_to_weight": 0.0006277755538141884, "update_to_weight": 4.689501902289656e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832523919556117, "grad_norm": 0.007453374979515504, "grad_rms": 9.048223588600356e-06, "grad_to_weight": 0.00024173742632817869, "update_to_weight": 3.929102051495923e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.635307017370295, "grad_norm": 0.003533857262152779, "grad_rms": 3.1723515714895678e-06, "grad_to_weight": 0.0001020304875709772, "update_to_weight": 4.182409579523526e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.086124295137324, "grad_norm": 0.003018334615899237, "grad_rms": 2.427833110253251e-06, "grad_to_weight": 8.138716766084388e-05, "update_to_weight": 4.5391819109943386e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.671997456551704, "grad_norm": 0.0011111254815266078, "grad_rms": 2.0074608161499735e-06, "grad_to_weight": 4.50358947824691e-05, "update_to_weight": 2.8384117662730465e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35952040174549, "grad_norm": 0.0031602097653543135, "grad_rms": 2.541952062987901e-06, "grad_to_weight": 8.458914170661206e-05, "update_to_weight": 4.985103327339476e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50484418038517, "grad_norm": 0.008137963066293412, "grad_rms": 3.61297363655093e-06, "grad_to_weight": 0.0001281471228112547, "update_to_weight": 5.6626032915783946e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78362254073034, "grad_norm": 0.018423209094640892, "grad_rms": 1.3746864493780666e-05, "grad_to_weight": 0.00036277855286642453, "update_to_weight": 4.832921464528828e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10856391612722, "grad_norm": 0.02706462625442797, "grad_rms": 2.324203286904298e-05, "grad_to_weight": 0.000674784225908063, "update_to_weight": 5.749589236072202e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20443426670153, "grad_norm": 0.07506705467963194, "grad_rms": 0.00011701890058784598, "grad_to_weight": 0.002661533784716148, "update_to_weight": 4.7565232639785484e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416447350190884, "grad_norm": 0.004361341899016411, "grad_rms": 0.00037126207744094826, "grad_to_weight": 0.0005181927382837483, "update_to_weight": 3.16220072690321e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3225749117541652, "grad_norm": 0.06646007221972969, "grad_rms": 0.0015385265287486214, "grad_to_weight": 0.050250516344349815, "update_to_weight": 8.034738589593558e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48650}
{"step": 48700, "lr": 1.0166487511031909e-05, "window_steps": 50, "elapsed_seconds": 4437.621481276, "steps_per_second": 1.735163765654463, "eta_seconds": 749.2088215141299, "images_seen": 49868800, "images_per_second": 1776.80769603017, "effective_batch_size": 1024, "loss": 0.366384357213974, "edge_loss": 0.19325007498264313, "diag_loss": 0.1731342226266861, "jvp_rms": 2.111523389816284, "grad_norm": 0.20781058073043823, "grad_norm_postclip": 0.20781058073043823, "clip_fraction": 0.0, "ema_updates": 48700, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4013878752787908, "updates_nfe_1": 18, "loss_nfe_3": 0.332863087952137, "updates_nfe_3": 10, "loss_nfe_2": 0.37716417014598846, "updates_nfe_2": 4, "loss_nfe_4": 0.3003696139369692, "updates_nfe_4": 7, "loss_nfe_5": 0.3776690418070013, "updates_nfe_5": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.509930021252078, "grad_norm": 0.06247153229883968, "grad_rms": 0.00018191100009904776, "grad_to_weight": 0.003783876262250927, "update_to_weight": 4.5299621654397285e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.260253488364178, "grad_norm": 0.03184883413741579, "grad_rms": 9.274069414867879e-05, "grad_to_weight": 0.0019586923512728005, "update_to_weight": 6.612445747689945e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066045202656736, "grad_norm": 0.0034406766611652082, "grad_rms": 6.15979933455185e-05, "grad_to_weight": 0.0006791642244646472, "update_to_weight": 2.5234660427061264e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.868362048889091, "grad_norm": 0.01453142468782631, "grad_rms": 4.650061852187399e-05, "grad_to_weight": 0.001129236544062015, "update_to_weight": 6.117240449887006e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.041738252157829, "grad_norm": 0.013318658092193191, "grad_rms": 0.00030300996349902544, "grad_to_weight": 0.0026416797989250824, "update_to_weight": 2.3711016361141626e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644792978609896, "grad_norm": 0.03326063224912971, "grad_rms": 8.014633448462327e-05, "grad_to_weight": 0.0019982604945506205, "update_to_weight": 6.368204571364655e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57489524035085, "grad_norm": 0.022997769275131866, "grad_rms": 3.172887918308693e-05, "grad_to_weight": 0.0009358236952876441, "update_to_weight": 8.246584287128158e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832495215769818, "grad_norm": 0.013929427004573796, "grad_rms": 1.6909999878560282e-05, "grad_to_weight": 0.000451777480450215, "update_to_weight": 7.21905460489544e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63365048754968, "grad_norm": 0.010064748549345167, "grad_rms": 9.035147293332332e-06, "grad_to_weight": 0.0002906060553149979, "update_to_weight": 8.628040611751153e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08528823557429, "grad_norm": 0.008237154429427115, "grad_rms": 6.625652488193176e-06, "grad_to_weight": 0.0002221138036491132, "update_to_weight": 8.895774774951343e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.672301739897136, "grad_norm": 0.0034124276788610455, "grad_rms": 6.1652036310492185e-06, "grad_to_weight": 0.00013831006587208157, "update_to_weight": 5.847183312728545e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.35981120353879, "grad_norm": 0.008216258410559379, "grad_rms": 6.608844528528157e-06, "grad_to_weight": 0.00021992237503012647, "update_to_weight": 8.885588801485576e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50491722472166, "grad_norm": 0.019683455677905196, "grad_rms": 8.738772326830087e-06, "grad_to_weight": 0.0003099516783598401, "update_to_weight": 9.131352353538476e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.783891758207886, "grad_norm": 0.030271659883484705, "grad_rms": 2.258783495765301e-05, "grad_to_weight": 0.0005960878309132755, "update_to_weight": 7.10290879614601e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.108999133903076, "grad_norm": 0.029797840544345065, "grad_rms": 2.5589209429591194e-05, "grad_to_weight": 0.0007429215684207323, "update_to_weight": 7.74621703241166e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204779426890244, "grad_norm": 0.04497661054746796, "grad_rms": 7.011216226471273e-05, "grad_to_weight": 0.0015946450020661238, "update_to_weight": 6.580201084983872e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416189450876619, "grad_norm": 0.0028976951160720873, "grad_rms": 0.00024666818917040944, "grad_to_weight": 0.00034430012929072875, "update_to_weight": 4.946872051654112e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3222770746769326, "grad_norm": 0.04274671255300161, "grad_rms": 0.0009895708668829998, "grad_to_weight": 0.03232810533559751, "update_to_weight": 9.756940613512572e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48700}
{"step": 48750, "lr": 1.0153943689205709e-05, "window_steps": 50, "elapsed_seconds": 4465.113667078999, "steps_per_second": 1.7356780986652725, "eta_seconds": 720.1796237224193, "images_seen": 49920000, "images_per_second": 1777.334373033239, "effective_batch_size": 1024, "loss": 0.36334049701690674, "edge_loss": 0.18867617845535278, "diag_loss": 0.17466428875923157, "jvp_rms": 1.9694530963897705, "grad_norm": 0.3385254740715027, "grad_norm_postclip": 0.3151838481426239, "clip_fraction": 0.03999999910593033, "ema_updates": 48750, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.34680306098677893, "updates_nfe_3": 11, "loss_nfe_4": 0.33700587451457975, "updates_nfe_4": 10, "loss_nfe_2": 0.35479937493801117, "updates_nfe_2": 14, "loss_nfe_1": 0.4004050601612438, "updates_nfe_1": 11, "loss_nfe_5": 0.40262099355459213, "updates_nfe_5": 4, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.512127667487974, "grad_norm": 0.08017145800347317, "grad_rms": 0.000233451454897025, "grad_to_weight": 0.004855307542305953, "update_to_weight": 6.779201074398761e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.257968298748924, "grad_norm": 0.04342667853725554, "grad_rms": 0.0001264542461661171, "grad_to_weight": 0.0026711011941508885, "update_to_weight": 5.7479395439883275e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066428838133918, "grad_norm": 0.003769208534168408, "grad_rms": 6.747965736685913e-05, "grad_to_weight": 0.0007439576582618486, "update_to_weight": 3.381135337887121e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.867349369997315, "grad_norm": 0.040823626155602054, "grad_rms": 0.00013063577091182034, "grad_to_weight": 0.003172652345228936, "update_to_weight": 7.944569530016252e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.04225689083992, "grad_norm": 0.023627223583501503, "grad_rms": 0.0005375379490983825, "grad_to_weight": 0.004685842886431312, "update_to_weight": 3.0043222834997955e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644717806540758, "grad_norm": 0.10644845345438904, "grad_rms": 0.000256503042155777, "grad_to_weight": 0.006395329418715573, "update_to_weight": 8.386718125931095e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574633215090923, "grad_norm": 0.02131693176783387, "grad_rms": 2.9409911218958008e-05, "grad_to_weight": 0.0008674364162938331, "update_to_weight": 8.367180334636172e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83266446874463, "grad_norm": 0.007769259870602077, "grad_rms": 9.431700487410268e-06, "grad_to_weight": 0.0002519814620133739, "update_to_weight": 6.311498743623314e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63405052913516, "grad_norm": 0.003694795974421464, "grad_rms": 3.316826613604893e-06, "grad_to_weight": 0.0001066810239626259, "update_to_weight": 6.782740196879911e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08594655176844, "grad_norm": 0.0032131923025092075, "grad_rms": 2.5845691927429303e-06, "grad_to_weight": 8.664177677179948e-05, "update_to_weight": 7.217792372261001e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673196014349255, "grad_norm": 0.0011769421047197193, "grad_rms": 2.126371141138605e-06, "grad_to_weight": 4.770124243471507e-05, "update_to_weight": 4.693605092206663e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36076691258994, "grad_norm": 0.0031806019277476117, "grad_rms": 2.5583547397445866e-06, "grad_to_weight": 8.513213701391669e-05, "update_to_weight": 7.568872796483831e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50622047534005, "grad_norm": 0.008590639378429002, "grad_rms": 3.813946234769161e-06, "grad_to_weight": 0.0001352724081850346, "update_to_weight": 8.512452723336857e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78423068175297, "grad_norm": 0.02216756128155137, "grad_rms": 1.654079153797964e-05, "grad_to_weight": 0.0004365048162384842, "update_to_weight": 7.401115614504941e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.1092380451499, "grad_norm": 0.04409563791351558, "grad_rms": 3.7867593519776405e-05, "grad_to_weight": 0.0010993885713779527, "update_to_weight": 9.57347201222648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204630110461807, "grad_norm": 0.12003738574698183, "grad_rms": 0.00018712127403291187, "grad_to_weight": 0.004255946108027736, "update_to_weight": 7.969134470693068e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41673606113745, "grad_norm": 0.006232960858047294, "grad_rms": 0.0005305848636376413, "grad_to_weight": 0.000740543699216934, "update_to_weight": 4.999554352699099e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3229396798042048, "grad_norm": 0.07698838616968418, "grad_rms": 0.0017822531720397268, "grad_to_weight": 0.058194933106155276, "update_to_weight": 0.00015448908292592693, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48750}
{"step": 48800, "lr": 1.0141890296666345e-05, "window_steps": 50, "elapsed_seconds": 4492.7488366, "steps_per_second": 1.7361308819352665, "eta_seconds": 691.1921287076923, "images_seen": 49971200, "images_per_second": 1777.798023101713, "effective_batch_size": 1024, "loss": 0.3604809045791626, "edge_loss": 0.18808603286743164, "diag_loss": 0.17239488661289215, "jvp_rms": 1.8731640577316284, "grad_norm": 0.42754065990448, "grad_norm_postclip": 0.40926283597946167, "clip_fraction": 0.05999999865889549, "ema_updates": 48800, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.30068184299902484, "updates_nfe_5": 11, "loss_nfe_4": 0.3661528428395589, "updates_nfe_4": 12, "loss_nfe_1": 0.4288741983473301, "updates_nfe_1": 8, "loss_nfe_2": 0.34919123351573944, "updates_nfe_2": 10, "loss_nfe_3": 0.3777560889720917, "updates_nfe_3": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.508140636528108, "grad_norm": 0.023615967248339927, "grad_rms": 6.876738991932274e-05, "grad_to_weight": 0.001430564941764798, "update_to_weight": 5.713200036587214e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26257870279044, "grad_norm": 0.015713868328578302, "grad_rms": 4.5757249708592856e-05, "grad_to_weight": 0.0009662593255202519, "update_to_weight": 4.677358290659295e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066676924511079, "grad_norm": 0.0010952564435341482, "grad_rms": 1.9608235752558344e-05, "grad_to_weight": 0.00021616859725861394, "update_to_weight": 2.7273313487460357e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.86907763515659, "grad_norm": 0.010860600887394675, "grad_rms": 3.4753967324769606e-05, "grad_to_weight": 0.0008439300154445392, "update_to_weight": 7.714799338090656e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.041349521887076, "grad_norm": 0.04125934521613288, "grad_rms": 0.000938682606114968, "grad_to_weight": 0.008184186602616019, "update_to_weight": 2.8885276967292402e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644661885634967, "grad_norm": 0.0901969724810514, "grad_rms": 0.00021734273334975032, "grad_to_weight": 0.005418972947650869, "update_to_weight": 9.671516550739374e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574776025413104, "grad_norm": 0.022444795685805447, "grad_rms": 3.0965969007005284e-05, "grad_to_weight": 0.0009133265614545168, "update_to_weight": 0.0001044757629020871, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.832972289024962, "grad_norm": 0.0098147397717089, "grad_rms": 1.1914865435111014e-05, "grad_to_weight": 0.00031831961186571913, "update_to_weight": 7.635872210819152e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63467789782283, "grad_norm": 0.006236638339747869, "grad_rms": 5.5986441925100835e-06, "grad_to_weight": 0.00018006918840552893, "update_to_weight": 8.062509955646188e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.086107837453774, "grad_norm": 0.005235993352870379, "grad_rms": 4.211633117217318e-06, "grad_to_weight": 0.0001411847631954108, "update_to_weight": 7.921408212585015e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673371649506834, "grad_norm": 0.0018478990305180092, "grad_rms": 3.3385832272244535e-06, "grad_to_weight": 7.489446747562548e-05, "update_to_weight": 5.095731055639204e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36129458420589, "grad_norm": 0.005269994435420262, "grad_rms": 4.238982289692838e-06, "grad_to_weight": 0.00014105492044828925, "update_to_weight": 8.021736299999552e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50665873060084, "grad_norm": 0.01361651740525761, "grad_rms": 6.04526194160278e-06, "grad_to_weight": 0.0002144108614345421, "update_to_weight": 8.844436183430176e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78449508862178, "grad_norm": 0.030392389060884196, "grad_rms": 2.2677919569602576e-05, "grad_to_weight": 0.0005984580334578059, "update_to_weight": 8.082488190453946e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.10969151608654, "grad_norm": 0.04565322589128067, "grad_rms": 3.920518860183373e-05, "grad_to_weight": 0.0011382093495526093, "update_to_weight": 0.00010594407045320933, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.204898446665585, "grad_norm": 0.12570131528178755, "grad_rms": 0.0001959505375493587, "grad_to_weight": 0.004456719300708849, "update_to_weight": 8.478026552988177e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416125541388862, "grad_norm": 0.022220287729940495, "grad_rms": 0.0018915165045450336, "grad_to_weight": 0.0026402039300228424, "update_to_weight": 5.529135970270857e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3224393879750629, "grad_norm": 0.28831175665983894, "grad_rms": 0.006674312430849227, "grad_to_weight": 0.218015100942589, "update_to_weight": 0.0002157097896915406, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48800}
{"step": 48850, "lr": 1.0130327465227797e-05, "window_steps": 50, "elapsed_seconds": 4520.241446614, "steps_per_second": 1.7366328973157483, "eta_seconds": 662.2009762555541, "images_seen": 50022400, "images_per_second": 1778.3120868513263, "effective_batch_size": 1024, "loss": 0.3755190968513489, "edge_loss": 0.19370850920677185, "diag_loss": 0.18181052803993225, "jvp_rms": 2.3349609375, "grad_norm": 0.29069364070892334, "grad_norm_postclip": 0.29069364070892334, "clip_fraction": 0.0, "ema_updates": 48850, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3982482850551605, "updates_nfe_3": 11, "loss_nfe_2": 0.3494185358285904, "updates_nfe_2": 10, "loss_nfe_4": 0.3528985335276677, "updates_nfe_4": 13, "loss_nfe_5": 0.38397620405469624, "updates_nfe_5": 7, "loss_nfe_1": 0.40283579958809745, "updates_nfe_1": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.511576897160953, "grad_norm": 0.27599746722116575, "grad_rms": 0.0008036776662822297, "grad_to_weight": 0.01671539120340599, "update_to_weight": 7.83555434525803e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.259724216351177, "grad_norm": 0.09180956103546481, "grad_rms": 0.00026734047416548116, "grad_to_weight": 0.0056464402356307414, "update_to_weight": 6.844165675737142e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067291629461226, "grad_norm": 0.008510903780552668, "grad_rms": 0.0001523696197193052, "grad_to_weight": 0.0016795764686347013, "update_to_weight": 3.2001987827104045e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.868031217594718, "grad_norm": 0.12806718295188926, "grad_rms": 0.00040981551001023414, "grad_to_weight": 0.009952352522799326, "update_to_weight": 0.00011166337194819954, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0420968789065626, "grad_norm": 0.04394538932369723, "grad_rms": 0.0009997922255192621, "grad_to_weight": 0.008715697135360734, "update_to_weight": 2.3375726290467406e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64481406339759, "grad_norm": 0.2310621574787287, "grad_rms": 0.0005567778995095153, "grad_to_weight": 0.013881930828343757, "update_to_weight": 7.0451812269191e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574460599767622, "grad_norm": 0.04424760743333184, "grad_rms": 6.10462246836675e-05, "grad_to_weight": 0.0018005525392386535, "update_to_weight": 8.148055507392487e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83347805779661, "grad_norm": 0.015990350969501514, "grad_rms": 1.941191355994908e-05, "grad_to_weight": 0.0005186035431853645, "update_to_weight": 5.970359887325044e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63603304589384, "grad_norm": 0.004491455323606578, "grad_rms": 4.031989493949886e-06, "grad_to_weight": 0.00012967580085326914, "update_to_weight": 5.913908984312113e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.086917546825056, "grad_norm": 0.0031352704189859617, "grad_rms": 2.521891805075403e-06, "grad_to_weight": 8.453844714992677e-05, "update_to_weight": 6.135378194549602e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673266515935953, "grad_norm": 0.0013071128508821872, "grad_rms": 2.3615495045860292e-06, "grad_to_weight": 5.297688694919864e-05, "update_to_weight": 3.9454584618371785e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36211139863568, "grad_norm": 0.0032630624183343027, "grad_rms": 2.624682809627737e-06, "grad_to_weight": 8.733613535699316e-05, "update_to_weight": 6.245336456487493e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.507344939999854, "grad_norm": 0.010574367655270037, "grad_rms": 4.694652857289099e-06, "grad_to_weight": 0.00016650621538753406, "update_to_weight": 7.032477281628296e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.784598447296254, "grad_norm": 0.038767437854336076, "grad_rms": 2.892713816669674e-05, "grad_to_weight": 0.0007633699790807352, "update_to_weight": 6.37571181073348e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.109997377816384, "grad_norm": 0.1138166608319563, "grad_rms": 9.774125632598565e-05, "grad_to_weight": 0.002837613270324092, "update_to_weight": 8.907076604788649e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20475995087919, "grad_norm": 0.3258973982619627, "grad_rms": 0.0005080278613808701, "grad_to_weight": 0.011554694981610858, "update_to_weight": 7.48543034773841e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416774585397112, "grad_norm": 0.007946994114920922, "grad_rms": 0.0006764930640227777, "grad_to_weight": 0.0009441852142159959, "update_to_weight": 6.207199819061967e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3230927668029726, "grad_norm": 0.12353810393544683, "grad_rms": 0.0028598622280697015, "grad_to_weight": 0.09337070463619533, "update_to_weight": 0.00012966887732069197, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48850}
{"step": 48900, "lr": 1.0119255321339333e-05, "window_steps": 50, "elapsed_seconds": 4547.717091981, "steps_per_second": 1.7371353231119167, "eta_seconds": 633.2264305289999, "images_seen": 50073600, "images_per_second": 1778.8265708666027, "effective_batch_size": 1024, "loss": 0.3742154538631439, "edge_loss": 0.19285233318805695, "diag_loss": 0.18136313557624817, "jvp_rms": 2.646601438522339, "grad_norm": 0.1628469079732895, "grad_norm_postclip": 0.1628469079732895, "clip_fraction": 0.0, "ema_updates": 48900, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.4054950140416622, "updates_nfe_1": 8, "loss_nfe_2": 0.381993060762232, "updates_nfe_2": 11, "loss_nfe_4": 0.33981708330767496, "updates_nfe_4": 7, "loss_nfe_3": 0.38204124569892883, "updates_nfe_3": 12, "loss_nfe_5": 0.3584730786581834, "updates_nfe_5": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.512452021954335, "grad_norm": 0.025280754857451477, "grad_rms": 7.36150888276399e-05, "grad_to_weight": 0.001531011555633236, "update_to_weight": 1.825521858228734e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.259805860498265, "grad_norm": 0.007914771540860914, "grad_rms": 2.304704164556297e-05, "grad_to_weight": 0.0004867691292728863, "update_to_weight": 2.244930470667181e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066817526495479, "grad_norm": 0.0010712643165472022, "grad_rms": 1.917870777767847e-05, "grad_to_weight": 0.00021142745144172464, "update_to_weight": 1.2780277522471276e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.868879622723208, "grad_norm": 0.007484864252810702, "grad_rms": 2.3951596267057086e-05, "grad_to_weight": 0.0005816251664670413, "update_to_weight": 2.5886133341333813e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.041868389707417, "grad_norm": 0.01992869558402456, "grad_rms": 0.000453393523559106, "grad_to_weight": 0.003952640974268081, "update_to_weight": 1.1548572418904527e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645289839114398, "grad_norm": 0.04973757888132054, "grad_rms": 0.00011984993561215165, "grad_to_weight": 0.0029880872824721446, "update_to_weight": 3.461739038126786e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.575009487456843, "grad_norm": 0.02007409327680847, "grad_rms": 2.7695228727187996e-05, "grad_to_weight": 0.000816849868849668, "update_to_weight": 4.577563134252244e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8331563858591, "grad_norm": 0.012934143848524203, "grad_rms": 1.5701749313594084e-05, "grad_to_weight": 0.0004194881538127619, "update_to_weight": 4.0256550724836796e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63531406626407, "grad_norm": 0.008967328956225967, "grad_rms": 8.049991269035438e-06, "grad_to_weight": 0.0002589071067486129, "update_to_weight": 4.620413835847649e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.086489653831485, "grad_norm": 0.007462356188154066, "grad_rms": 6.002434368498957e-06, "grad_to_weight": 0.00020121495072217258, "update_to_weight": 5.02954867237095e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673136644435786, "grad_norm": 0.0025802139340492564, "grad_rms": 4.661650242033458e-06, "grad_to_weight": 0.00010457583773124114, "update_to_weight": 3.136322518970539e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36160497579283, "grad_norm": 0.007572287695582328, "grad_rms": 6.090859075351704e-06, "grad_to_weight": 0.00020267565326726548, "update_to_weight": 5.4256297139428057e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.506947054780944, "grad_norm": 0.0182831376287699, "grad_rms": 8.117079631360947e-06, "grad_to_weight": 0.0002878919311457203, "update_to_weight": 6.33832913479911e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78461767679388, "grad_norm": 0.026078407749402538, "grad_rms": 1.9458951787551438e-05, "grad_to_weight": 0.0005135099749174466, "update_to_weight": 4.860131551693798e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.110428015004715, "grad_norm": 0.022958361803016317, "grad_rms": 1.971573501990575e-05, "grad_to_weight": 0.0005723788784908487, "update_to_weight": 4.936121418023008e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20530743806529, "grad_norm": 0.03764389495503835, "grad_rms": 5.868149776599904e-05, "grad_to_weight": 0.0013346387036446531, "update_to_weight": 3.295609432675599e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416575163793938, "grad_norm": 0.00395992058583184, "grad_rms": 0.0003370908259975364, "grad_to_weight": 0.0004704907291586317, "update_to_weight": 1.4139753467864484e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3228254161612374, "grad_norm": 0.03632643731616756, "grad_rms": 0.0008409438274618081, "grad_to_weight": 0.027461248379687754, "update_to_weight": 2.674033802725478e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48900}
{"step": 48950, "lr": 1.0108673986084137e-05, "window_steps": 50, "elapsed_seconds": 4575.280669598, "steps_per_second": 1.7375983189023714, "eta_seconds": 604.282352588415, "images_seen": 50124800, "images_per_second": 1779.3006785560283, "effective_batch_size": 1024, "loss": 0.34007614850997925, "edge_loss": 0.17576217651367188, "diag_loss": 0.16431394219398499, "jvp_rms": 1.6454296112060547, "grad_norm": 0.14636296033859253, "grad_norm_postclip": 0.14636296033859253, "clip_fraction": 0.0, "ema_updates": 48950, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.34182558125919765, "updates_nfe_4": 9, "loss_nfe_5": 0.31312815348307294, "updates_nfe_5": 12, "loss_nfe_1": 0.39370829463005064, "updates_nfe_1": 10, "loss_nfe_2": 0.35326617104666574, "updates_nfe_2": 7, "loss_nfe_3": 0.31332431733608246, "updates_nfe_3": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.509806175457417, "grad_norm": 0.028239801670144105, "grad_rms": 8.223154411901768e-05, "grad_to_weight": 0.0017104865659854846, "update_to_weight": 2.5923582691102094e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.263423129365037, "grad_norm": 0.013639295238650394, "grad_rms": 3.9716295506251844e-05, "grad_to_weight": 0.0008386484893222417, "update_to_weight": 3.438257796545135e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066832772259315, "grad_norm": 0.0008385581041455756, "grad_rms": 1.5012598278123863e-05, "grad_to_weight": 0.00016549946324193767, "update_to_weight": 1.7863410562689875e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.86913454834619, "grad_norm": 0.00942068665279159, "grad_rms": 3.01462358761397e-05, "grad_to_weight": 0.0007320373112426771, "update_to_weight": 4.3208209729784284e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.041410245233191, "grad_norm": 0.016131402612472402, "grad_rms": 0.0003670021171020571, "grad_to_weight": 0.003199779789340719, "update_to_weight": 1.5228537983994922e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644760893012158, "grad_norm": 0.08182068605263941, "grad_rms": 0.00019715885203318027, "grad_to_weight": 0.004915702098609873, "update_to_weight": 4.7053013671229294e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57476981628592, "grad_norm": 0.019451151457127645, "grad_rms": 2.6835786861401563e-05, "grad_to_weight": 0.0007915089989667856, "update_to_weight": 6.651259468053678e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833631469384915, "grad_norm": 0.0077845868259142515, "grad_rms": 9.450307054096927e-06, "grad_to_weight": 0.0002524706450371784, "update_to_weight": 5.8804445248454894e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63650883303087, "grad_norm": 0.004498437169986109, "grad_rms": 4.038257113067164e-06, "grad_to_weight": 0.0001298755943236463, "update_to_weight": 6.840891573779262e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08777990078066, "grad_norm": 0.003933158919064753, "grad_rms": 3.1636828472539625e-06, "grad_to_weight": 0.00010604999624099808, "update_to_weight": 6.986258494208858e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67280515979334, "grad_norm": 0.0013127183203251226, "grad_rms": 2.3716768578417194e-06, "grad_to_weight": 5.32050697852679e-05, "update_to_weight": 4.549762441003746e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.362702761138515, "grad_norm": 0.003991527491704861, "grad_rms": 3.210632298288124e-06, "grad_to_weight": 0.00010683187234132607, "update_to_weight": 7.004680444877589e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50865006290482, "grad_norm": 0.010599432387310053, "grad_rms": 4.7057807298697535e-06, "grad_to_weight": 0.0001668974600595572, "update_to_weight": 7.735288607350714e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78507918255248, "grad_norm": 0.021971105812260902, "grad_rms": 1.63942021625063e-05, "grad_to_weight": 0.00043262915340317533, "update_to_weight": 5.634852654166357e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11053148897587, "grad_norm": 0.031073960098482484, "grad_rms": 2.6685090538137364e-05, "grad_to_weight": 0.000774708261021746, "update_to_weight": 6.412389361001777e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20498500618674, "grad_norm": 0.05955609167486086, "grad_rms": 9.283950730242444e-05, "grad_to_weight": 0.002111544879807498, "update_to_weight": 4.3963063767771946e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416680313955391, "grad_norm": 0.006573992156735876, "grad_rms": 0.0005596153756578206, "grad_to_weight": 0.0007810671085886177, "update_to_weight": 2.803291008221864e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3229785163818966, "grad_norm": 0.0856080350462904, "grad_rms": 0.0019817949122489715, "grad_to_weight": 0.06470856025720861, "update_to_weight": 5.0430633159142944e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 48950}
{"step": 49000, "lr": 1.009858357517801e-05, "window_steps": 50, "elapsed_seconds": 4603.151212176, "steps_per_second": 1.7379398658116736, "eta_seconds": 575.393901522, "images_seen": 50176000, "images_per_second": 1779.6504225911538, "effective_batch_size": 1024, "loss": 0.37687477469444275, "edge_loss": 0.19369304180145264, "diag_loss": 0.18318170309066772, "jvp_rms": 2.571601390838623, "grad_norm": 0.18482986092567444, "grad_norm_postclip": 0.18482986092567444, "clip_fraction": 0.0, "ema_updates": 49000, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3462676964700222, "updates_nfe_3": 8, "loss_nfe_2": 0.3740378374403173, "updates_nfe_2": 11, "loss_nfe_1": 0.3976626917719841, "updates_nfe_1": 8, "loss_nfe_4": 0.4342987338701884, "updates_nfe_4": 12, "loss_nfe_5": 0.32420841130343353, "updates_nfe_5": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51257215162132, "grad_norm": 0.08898224037664501, "grad_rms": 0.0002591075925677272, "grad_to_weight": 0.005388757097294991, "update_to_weight": 4.127103296619732e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.25917991156025, "grad_norm": 0.07958610209931724, "grad_rms": 0.00023174695567921303, "grad_to_weight": 0.004894841100979002, "update_to_weight": 4.7284653166442384e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066582058317235, "grad_norm": 0.0037163079040714907, "grad_rms": 6.65325841653976e-05, "grad_to_weight": 0.0007334940718014915, "update_to_weight": 2.3897320716767852e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.868605127394526, "grad_norm": 0.033680407829762254, "grad_rms": 0.00010777744301045456, "grad_to_weight": 0.0026172539678029146, "update_to_weight": 6.089566204171855e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.041637998557224, "grad_norm": 0.015561575332168455, "grad_rms": 0.00035403809758819516, "grad_to_weight": 0.0030866110055148234, "update_to_weight": 2.1109965509831923e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64514866623587, "grad_norm": 0.08233529182710825, "grad_rms": 0.000198398869547561, "grad_to_weight": 0.004946503841934596, "update_to_weight": 6.451014512658308e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574605894760804, "grad_norm": 0.017148482674558742, "grad_rms": 2.365890919440985e-05, "grad_to_weight": 0.0006978131306762776, "update_to_weight": 7.362120811979275e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833765085307657, "grad_norm": 0.007801370684328521, "grad_rms": 9.470682267208988e-06, "grad_to_weight": 0.0002530138847054357, "update_to_weight": 6.124072897483363e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63699342432253, "grad_norm": 0.0038313585255421068, "grad_rms": 3.4394191213142398e-06, "grad_to_weight": 0.0001106146390538527, "update_to_weight": 6.230124492163718e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.087615330963246, "grad_norm": 0.0032094286442793368, "grad_rms": 2.581541849777699e-06, "grad_to_weight": 8.653639808434621e-05, "update_to_weight": 6.5535772525455e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.672765579239897, "grad_norm": 0.0012183404480142217, "grad_rms": 2.201165170614968e-06, "grad_to_weight": 4.9379970968449316e-05, "update_to_weight": 4.410843214874663e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.362838348459626, "grad_norm": 0.0033993050459884644, "grad_rms": 2.7342711769029577e-06, "grad_to_weight": 9.098091034426481e-05, "update_to_weight": 6.820037763861075e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.508750012301455, "grad_norm": 0.008547712619641122, "grad_rms": 3.794888241197579e-06, "grad_to_weight": 0.0001345911015094055, "update_to_weight": 7.542528684226318e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78516811807607, "grad_norm": 0.020402998102508927, "grad_rms": 1.5224125652660696e-05, "grad_to_weight": 0.0004017511186547957, "update_to_weight": 6.234114536580851e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.110666917740694, "grad_norm": 0.04399183777467068, "grad_rms": 3.777845405721133e-05, "grad_to_weight": 0.001096761563822649, "update_to_weight": 7.711810524732254e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20515812443211, "grad_norm": 0.11128842750306565, "grad_rms": 0.00017348288793449164, "grad_to_weight": 0.003945676425996154, "update_to_weight": 6.56815762837317e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416986692280519, "grad_norm": 0.004173665185717484, "grad_rms": 0.0003552859746542434, "grad_to_weight": 0.0004958621580743715, "update_to_weight": 5.3863947954885615e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3232198452910868, "grad_norm": 0.06693149158169874, "grad_rms": 0.001549439715724362, "grad_to_weight": 0.05058229123443573, "update_to_weight": 0.0001488468031296789, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49000}
{"step": 49050, "lr": 1.0088984198968057e-05, "window_steps": 50, "elapsed_seconds": 4653.982220166, "steps_per_second": 1.729701494156733, "eta_seconds": 549.2277154233168, "images_seen": 50227200, "images_per_second": 1771.2143300164946, "effective_batch_size": 1024, "loss": 0.37406042218208313, "edge_loss": 0.19561994075775146, "diag_loss": 0.17844045162200928, "jvp_rms": 2.3907811641693115, "grad_norm": 0.1935448795557022, "grad_norm_postclip": 0.1935448795557022, "clip_fraction": 0.0, "ema_updates": 49050, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.3703906481916254, "updates_nfe_4": 11, "loss_nfe_1": 0.4015184428010668, "updates_nfe_1": 14, "loss_nfe_3": 0.35965273529291153, "updates_nfe_3": 10, "loss_nfe_5": 0.34178973899947274, "updates_nfe_5": 9, "loss_nfe_2": 0.38913855453332263, "updates_nfe_2": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51155471810527, "grad_norm": 0.07170150279348304, "grad_rms": 0.00020878777263495148, "grad_to_weight": 0.00434250462888639, "update_to_weight": 4.063987203151289e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26229064947398, "grad_norm": 0.08256229725786825, "grad_rms": 0.00024041334527875248, "grad_to_weight": 0.005076916840158604, "update_to_weight": 5.6099454360342434e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066475896422799, "grad_norm": 0.003446495855244177, "grad_rms": 6.170217363139028e-05, "grad_to_weight": 0.0006802550580922701, "update_to_weight": 2.5658601644573287e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.869647941368312, "grad_norm": 0.03308462285579255, "grad_rms": 0.00010587092865341155, "grad_to_weight": 0.0025707480893432245, "update_to_weight": 5.380775916065667e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.04129787809194, "grad_norm": 0.01991417539691014, "grad_rms": 0.0004530631778638352, "grad_to_weight": 0.003950208037388851, "update_to_weight": 2.1559855547881117e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645097330346818, "grad_norm": 0.11223827554541237, "grad_rms": 0.00027045446119188866, "grad_to_weight": 0.006743023084688311, "update_to_weight": 6.384197369306803e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57559064929506, "grad_norm": 0.016102747653070765, "grad_rms": 2.2216160562689587e-05, "grad_to_weight": 0.0006552333932829674, "update_to_weight": 7.415097698830235e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.834023407783143, "grad_norm": 0.00824709008682008, "grad_rms": 1.0011775238192601e-05, "grad_to_weight": 0.0002674672058768154, "update_to_weight": 6.556984841148207e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.635705277619586, "grad_norm": 0.003682055271124341, "grad_rms": 3.3053892557468427e-06, "grad_to_weight": 0.00010630807837204804, "update_to_weight": 7.700845798014131e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.087134782916976, "grad_norm": 0.0030915912099763795, "grad_rms": 2.486757917233898e-06, "grad_to_weight": 8.336020639158202e-05, "update_to_weight": 7.917348861975973e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67302285170232, "grad_norm": 0.00113111493012705, "grad_rms": 2.043575580385871e-06, "grad_to_weight": 4.584419740238714e-05, "update_to_weight": 5.4125879273586335e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36353261317695, "grad_norm": 0.00322547514092065, "grad_rms": 2.5944490389422576e-06, "grad_to_weight": 8.63268249903952e-05, "update_to_weight": 7.914246819647017e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.508292550313854, "grad_norm": 0.008553756676562597, "grad_rms": 3.797571593055681e-06, "grad_to_weight": 0.00013468724056446584, "update_to_weight": 8.618335925126446e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.7850287055644, "grad_norm": 0.019957823114977305, "grad_rms": 1.4891948983646075e-05, "grad_to_weight": 0.0003929863509713951, "update_to_weight": 6.54370101187266e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.111236017889205, "grad_norm": 0.05527851153899065, "grad_rms": 4.747100403541661e-05, "grad_to_weight": 0.0013781303451815097, "update_to_weight": 7.664662471613715e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205147304572904, "grad_norm": 0.12163511902016805, "grad_rms": 0.0001896119137930409, "grad_to_weight": 0.004312514935897794, "update_to_weight": 5.78959839312165e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.41681718849071, "grad_norm": 0.005387387838098124, "grad_rms": 0.00045860490809112843, "grad_to_weight": 0.0006400742367869086, "update_to_weight": 2.484088648229448e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3232480883026871, "grad_norm": 0.0842233849773424, "grad_rms": 0.001949740766158566, "grad_to_weight": 0.06364897536740417, "update_to_weight": 8.033563966701227e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49050}
{"step": 49100, "lr": 1.0079875962431528e-05, "window_steps": 50, "elapsed_seconds": 4681.449003958, "steps_per_second": 1.7302335223884178, "eta_seconds": 520.1610004397778, "images_seen": 50278400, "images_per_second": 1771.7591269257398, "effective_batch_size": 1024, "loss": 0.3725816309452057, "edge_loss": 0.19327567517757416, "diag_loss": 0.17930586636066437, "jvp_rms": 2.0662498474121094, "grad_norm": 0.38181272149086, "grad_norm_postclip": 0.3492545187473297, "clip_fraction": 0.03999999910593033, "ema_updates": 49100, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3790635558275076, "updates_nfe_3": 13, "loss_nfe_1": 0.40064442970535974, "updates_nfe_1": 11, "loss_nfe_2": 0.34694780815731396, "updates_nfe_2": 11, "loss_nfe_5": 0.33960455159346264, "updates_nfe_5": 9, "loss_nfe_4": 0.40354955196380615, "updates_nfe_4": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51404413111619, "grad_norm": 0.0959513180486027, "grad_rms": 0.0002794008660384231, "grad_to_weight": 0.005810285917052186, "update_to_weight": 7.666698331153827e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26073582409528, "grad_norm": 0.14479854070574544, "grad_rms": 0.00042163920722581774, "grad_to_weight": 0.008904796330998863, "update_to_weight": 0.00010019947076848273, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066672783563502, "grad_norm": 0.004818515149390204, "grad_rms": 8.626525923156663e-05, "grad_to_weight": 0.0009510215787018393, "update_to_weight": 3.915024871730273e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.868752750879988, "grad_norm": 0.057835363847947685, "grad_rms": 0.00018507340120753774, "grad_to_weight": 0.004494247808436043, "update_to_weight": 9.435245650408515e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.041777218113653, "grad_norm": 0.03264262640748693, "grad_rms": 0.0007426454653147548, "grad_to_weight": 0.006474428558685888, "update_to_weight": 2.7468635601898074e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645236670248753, "grad_norm": 0.14306865859618476, "grad_rms": 0.000344744756510641, "grad_to_weight": 0.008595171185033482, "update_to_weight": 8.287656231457851e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.574884063975002, "grad_norm": 0.023731892399802872, "grad_rms": 3.274171236919006e-05, "grad_to_weight": 0.0009656970237589892, "update_to_weight": 8.788312322221204e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83396303375853, "grad_norm": 0.009741554173538926, "grad_rms": 1.1826019823890952e-05, "grad_to_weight": 0.00031593584525198385, "update_to_weight": 6.653613726451795e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63618283143818, "grad_norm": 0.003178663164384454, "grad_rms": 2.8534930351510816e-06, "grad_to_weight": 9.177290638099072e-05, "update_to_weight": 7.037415818878774e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08772723851854, "grad_norm": 0.002687524868379795, "grad_rms": 2.1617423812825197e-06, "grad_to_weight": 7.246399465504556e-05, "update_to_weight": 7.084539468491903e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673439676872807, "grad_norm": 0.0009934225018123624, "grad_rms": 1.7948078587217941e-06, "grad_to_weight": 4.0262829780621495e-05, "update_to_weight": 4.624439340176125e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36391486165579, "grad_norm": 0.0027441514870500338, "grad_rms": 2.2072906710596398e-06, "grad_to_weight": 7.344389626222442e-05, "update_to_weight": 7.334769219495072e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.50977448454785, "grad_norm": 0.008146881058560679, "grad_rms": 3.616932916126191e-06, "grad_to_weight": 0.00012827759387718885, "update_to_weight": 8.391266471635423e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78534598865602, "grad_norm": 0.029168811685337213, "grad_rms": 2.1764921606387092e-05, "grad_to_weight": 0.0005743548875664387, "update_to_weight": 7.428984808759443e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11150534856707, "grad_norm": 0.0719444250589964, "grad_rms": 6.178303281361664e-05, "grad_to_weight": 0.0017936106968263288, "update_to_weight": 9.579258017395801e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20502287589371, "grad_norm": 0.20505103886068368, "grad_rms": 0.00031964551205954594, "grad_to_weight": 0.007270018526946023, "update_to_weight": 8.292570145556145e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417485666242994, "grad_norm": 0.007301756070134138, "grad_rms": 0.000621566754070902, "grad_to_weight": 0.0008674509657220665, "update_to_weight": 6.864858168938847e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.323931457502861, "grad_norm": 0.10338047489811783, "grad_rms": 0.0023932204385743614, "grad_to_weight": 0.07808597213416876, "update_to_weight": 0.0002570221397746701, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49100}
{"step": 49150, "lr": 1.007125896517463e-05, "window_steps": 50, "elapsed_seconds": 4708.931535096, "steps_per_second": 1.7307535561427203, "eta_seconds": 491.1155588750429, "images_seen": 50329600, "images_per_second": 1772.2916414901456, "effective_batch_size": 1024, "loss": 0.37709343433380127, "edge_loss": 0.19523324072360992, "diag_loss": 0.18186011910438538, "jvp_rms": 2.284374952316284, "grad_norm": 0.2766250669956207, "grad_norm_postclip": 0.2766250669956207, "clip_fraction": 0.0, "ema_updates": 49150, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.3580601703036915, "updates_nfe_2": 11, "loss_nfe_3": 0.34479056149721143, "updates_nfe_3": 10, "loss_nfe_4": 0.4230850802527534, "updates_nfe_4": 9, "loss_nfe_5": 0.3642428174614906, "updates_nfe_5": 10, "loss_nfe_1": 0.40179102718830106, "updates_nfe_1": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51102333598867, "grad_norm": 0.0272382954919248, "grad_rms": 7.931525595093263e-05, "grad_to_weight": 0.0016497036517751239, "update_to_weight": 3.6060111212091684e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.263476608238435, "grad_norm": 0.0212984629752347, "grad_rms": 6.201904384006036e-05, "grad_to_weight": 0.001309588563889577, "update_to_weight": 5.141898909254829e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06693083342162, "grad_norm": 0.0021620154325256774, "grad_rms": 3.870628522836089e-05, "grad_to_weight": 0.0004266913253018893, "update_to_weight": 2.7857804995696527e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.869822822305638, "grad_norm": 0.012255354838605764, "grad_rms": 3.921718568156824e-05, "grad_to_weight": 0.0009522551326320597, "update_to_weight": 5.442318708695683e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.04088887081536, "grad_norm": 0.03646708529590544, "grad_rms": 0.0008296549177807269, "grad_to_weight": 0.007234256939690662, "update_to_weight": 1.7212925973750507e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64496532304795, "grad_norm": 0.07473247070007827, "grad_rms": 0.0001800787654524367, "grad_to_weight": 0.004489794316158636, "update_to_weight": 5.7588710989410614e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.575574506096242, "grad_norm": 0.03355170220283488, "grad_rms": 4.6289615868604944e-05, "grad_to_weight": 0.001365245894638679, "update_to_weight": 7.861979504247924e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.834168898799387, "grad_norm": 0.019264079178390535, "grad_rms": 2.3386143339578592e-05, "grad_to_weight": 0.000624764015583395, "update_to_weight": 6.931001442443121e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.637090341767305, "grad_norm": 0.013334591724937358, "grad_rms": 1.1970492828566309e-05, "grad_to_weight": 0.0003849801352643578, "update_to_weight": 7.985038266782187e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08876401304882, "grad_norm": 0.010882073699782929, "grad_rms": 8.75312455331529e-06, "grad_to_weight": 0.0002934062104618618, "update_to_weight": 8.386900621694008e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67341493967052, "grad_norm": 0.004142016489660335, "grad_rms": 7.483345437651246e-06, "grad_to_weight": 0.00016787366077164698, "update_to_weight": 5.502819379569613e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3644751589829, "grad_norm": 0.010450948506306417, "grad_rms": 8.406343910150988e-06, "grad_to_weight": 0.00027970280491933724, "update_to_weight": 8.015299967120069e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51071821452659, "grad_norm": 0.021941815030359944, "grad_rms": 9.741405631480065e-06, "grad_to_weight": 0.0003454820799891579, "update_to_weight": 8.763440038623298e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78579787327359, "grad_norm": 0.030105217662109197, "grad_rms": 2.246364059761839e-05, "grad_to_weight": 0.0005927881203566224, "update_to_weight": 6.791265952966556e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11156469237161, "grad_norm": 0.029342070130906155, "grad_rms": 2.5197811786398122e-05, "grad_to_weight": 0.0007315114819364404, "update_to_weight": 7.24584926863298e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20542320968611, "grad_norm": 0.06433443860413914, "grad_rms": 0.0001002882730316488, "grad_to_weight": 0.0022809244210186446, "update_to_weight": 6.0402886832012314e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.416923695281282, "grad_norm": 0.006929226428773438, "grad_rms": 0.000589854925059951, "grad_to_weight": 0.000823249286750469, "update_to_weight": 4.436598981624776e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3233394798180906, "grad_norm": 0.0747116311844986, "grad_rms": 0.0017295471212158938, "grad_to_weight": 0.05645688980333953, "update_to_weight": 0.00013142045956984056, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49150}
{"step": 49200, "lr": 1.0063133301431472e-05, "window_steps": 50, "elapsed_seconds": 4736.417098456, "steps_per_second": 1.7312664466719108, "eta_seconds": 462.0894730200975, "images_seen": 50380800, "images_per_second": 1772.8168413920366, "effective_batch_size": 1024, "loss": 0.35469573736190796, "edge_loss": 0.18345864117145538, "diag_loss": 0.1712370663881302, "jvp_rms": 1.8716405630111694, "grad_norm": 0.20172210037708282, "grad_norm_postclip": 0.20172210037708282, "clip_fraction": 0.0, "ema_updates": 49200, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.33941249286427216, "updates_nfe_3": 17, "loss_nfe_5": 0.3433689493685961, "updates_nfe_5": 8, "loss_nfe_2": 0.3662428344999041, "updates_nfe_2": 7, "loss_nfe_1": 0.4008969455957413, "updates_nfe_1": 10, "loss_nfe_4": 0.3306440971791744, "updates_nfe_4": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.513271658552117, "grad_norm": 0.00422074360027637, "grad_rms": 1.2290393099613322e-05, "grad_to_weight": 0.00025559705475386365, "update_to_weight": 4.548223210541038e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.259926918306608, "grad_norm": 0.0065204799011125, "grad_rms": 1.8987000579128518e-05, "grad_to_weight": 0.0004010153264447499, "update_to_weight": 4.879841606664977e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067096836668197, "grad_norm": 0.0010079938836455474, "grad_rms": 1.804598532548375e-05, "grad_to_weight": 0.0001989292717579758, "update_to_weight": 2.477097674655053e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.869605258202489, "grad_norm": 0.0048808015467271745, "grad_rms": 1.5618584941328476e-05, "grad_to_weight": 0.0003792502915826714, "update_to_weight": 6.21937012280579e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.041144078601188, "grad_norm": 0.016657444223360716, "grad_rms": 0.00037896997814412917, "grad_to_weight": 0.0033042983822003376, "update_to_weight": 2.055066886051206e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645347591310518, "grad_norm": 0.04634358759952573, "grad_rms": 0.00011167161962371342, "grad_to_weight": 0.00278417662024185, "update_to_weight": 6.417098730348567e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57596193994118, "grad_norm": 0.020741660707235453, "grad_rms": 2.8616238330040448e-05, "grad_to_weight": 0.0008439816418142245, "update_to_weight": 9.488314154785482e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.834186713978692, "grad_norm": 0.01278688571900978, "grad_rms": 1.5522981374942463e-05, "grad_to_weight": 0.00041469832941022053, "update_to_weight": 8.590040051054784e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.637129108669285, "grad_norm": 0.009349743622124243, "grad_rms": 8.393285770292125e-06, "grad_to_weight": 0.00026993413896373147, "update_to_weight": 0.00010324929702731534, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.08857969969496, "grad_norm": 0.007537844210018127, "grad_rms": 6.063154050784577e-06, "grad_to_weight": 0.00020323895579317972, "update_to_weight": 0.00010267470991916183, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673438440013282, "grad_norm": 0.0025190020396039844, "grad_rms": 4.551059240725161e-06, "grad_to_weight": 0.00010209367639327, "update_to_weight": 6.877998774069787e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.364382048930466, "grad_norm": 0.00761853136112215, "grad_rms": 6.128055714102642e-06, "grad_to_weight": 0.00020389822990101415, "update_to_weight": 9.74058851354689e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51123524190917, "grad_norm": 0.018453253316493702, "grad_rms": 8.19260510252652e-06, "grad_to_weight": 0.00029055100638818865, "update_to_weight": 0.00010212622581284029, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78574739699983, "grad_norm": 0.02453758107366361, "grad_rms": 1.8309231594344326e-05, "grad_to_weight": 0.0004831588059903824, "update_to_weight": 7.275292854391476e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11196031549147, "grad_norm": 0.02333074710656453, "grad_rms": 2.0035524821680682e-05, "grad_to_weight": 0.0005816406608667804, "update_to_weight": 7.896955735457023e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205691538344702, "grad_norm": 0.04213922285866798, "grad_rms": 6.568907694050655e-05, "grad_to_weight": 0.001493997153070369, "update_to_weight": 5.507840178925467e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417038358532288, "grad_norm": 0.00645888405176337, "grad_rms": 0.0005498167230477097, "grad_to_weight": 0.0007673582769426313, "update_to_weight": 3.2178362092381597e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.323518235807668, "grad_norm": 0.08189350438673966, "grad_rms": 0.0018958048768684143, "grad_to_weight": 0.06187561468449636, "update_to_weight": 8.309855548459556e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49200}
{"step": 49250, "lr": 1.0055499060063002e-05, "window_steps": 50, "elapsed_seconds": 4763.892243742999, "steps_per_second": 1.7317772060935532, "eta_seconds": 433.0811130675454, "images_seen": 50432000, "images_per_second": 1773.3398590397985, "effective_batch_size": 1024, "loss": 0.36100471019744873, "edge_loss": 0.18633292615413666, "diag_loss": 0.17467175424098969, "jvp_rms": 1.9913281202316284, "grad_norm": 0.2046109139919281, "grad_norm_postclip": 0.2046109139919281, "clip_fraction": 0.0, "ema_updates": 49250, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.34807750284671785, "updates_nfe_3": 10, "loss_nfe_4": 0.36437993428923865, "updates_nfe_4": 11, "loss_nfe_1": 0.3970090217060513, "updates_nfe_1": 9, "loss_nfe_5": 0.34613877128471027, "updates_nfe_5": 11, "loss_nfe_2": 0.3534082273642222, "updates_nfe_2": 9, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.511637889410476, "grad_norm": 0.05343000679085257, "grad_rms": 0.0001555829609577774, "grad_to_weight": 0.003235899863399936, "update_to_weight": 5.174240846826766e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262309415317716, "grad_norm": 0.018177393922684488, "grad_rms": 5.2930795611864717e-05, "grad_to_weight": 0.0011177621491793117, "update_to_weight": 7.507714933188586e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06715838070765, "grad_norm": 0.0027328333363005274, "grad_rms": 4.8925565009890386e-05, "grad_to_weight": 0.0005393226599557907, "update_to_weight": 2.6517148687973504e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.869989994403829, "grad_norm": 0.012359443710400797, "grad_rms": 3.955027049766118e-05, "grad_to_weight": 0.0009603304832229839, "update_to_weight": 5.659234903097498e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040869952006982, "grad_norm": 0.025870230633058985, "grad_rms": 0.000588568126426296, "grad_to_weight": 0.005132096419737818, "update_to_weight": 2.1626575126251554e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.644891985207316, "grad_norm": 0.05498633946231759, "grad_rms": 0.00013249758818842828, "grad_to_weight": 0.0033034963225465905, "update_to_weight": 4.0360026034143706e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.575503724253167, "grad_norm": 0.026837676830368806, "grad_rms": 3.702660877749938e-05, "grad_to_weight": 0.0010920499181419885, "update_to_weight": 6.344584261378509e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83436486520559, "grad_norm": 0.016851931954553125, "grad_rms": 2.0457852804098075e-05, "grad_to_weight": 0.0005465308602341066, "update_to_weight": 5.147357195483316e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63786743365652, "grad_norm": 0.013201055728731865, "grad_rms": 1.1850617265976422e-05, "grad_to_weight": 0.000381116295742411, "update_to_weight": 6.5370809020303e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.089280743818016, "grad_norm": 0.011070936958652524, "grad_rms": 8.905038947027178e-06, "grad_to_weight": 0.0002984942478426954, "update_to_weight": 6.93577063693217e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67377733720432, "grad_norm": 0.004420236733394204, "grad_rms": 7.9860035503858e-06, "grad_to_weight": 0.00017914714366531776, "update_to_weight": 4.633739999498585e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.364826361478485, "grad_norm": 0.011377626049212866, "grad_rms": 9.15172794058465e-06, "grad_to_weight": 0.0003045009747708263, "update_to_weight": 7.404517121603815e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51256911467899, "grad_norm": 0.021785194698841544, "grad_rms": 9.671871630881366e-06, "grad_to_weight": 0.0003430060380569074, "update_to_weight": 8.621760233380876e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78588200028503, "grad_norm": 0.029289369813894676, "grad_rms": 2.1854878586649845e-05, "grad_to_weight": 0.0005767226768598858, "update_to_weight": 6.437396819759472e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.112193122657544, "grad_norm": 0.030397657333723915, "grad_rms": 2.6104308415370323e-05, "grad_to_weight": 0.0007578158900654492, "update_to_weight": 6.881548453647386e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20534097942489, "grad_norm": 0.052077525033536416, "grad_rms": 8.118148168685201e-05, "grad_to_weight": 0.0018463710497783276, "update_to_weight": 5.1380052812954205e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417162990729121, "grad_norm": 0.00543553183485115, "grad_rms": 0.00046270319725641256, "grad_to_weight": 0.0006457676821558504, "update_to_weight": 3.849296792034292e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3238226374075968, "grad_norm": 0.07131472217238577, "grad_rms": 0.0016509099115902116, "grad_to_weight": 0.05387029965890243, "update_to_weight": 7.983182288233713e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49250}
{"step": 49300, "lr": 1.0048356324556068e-05, "window_steps": 50, "elapsed_seconds": 4791.363785043, "steps_per_second": 1.7322834108129639, "eta_seconds": 404.09092163013247, "images_seen": 50483200, "images_per_second": 1773.858212672475, "effective_batch_size": 1024, "loss": 0.36203086376190186, "edge_loss": 0.18648776412010193, "diag_loss": 0.17554309964179993, "jvp_rms": 2.215703010559082, "grad_norm": 0.16771213710308075, "grad_norm_postclip": 0.16771213710308075, "clip_fraction": 0.0, "ema_updates": 49300, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3272811050216357, "updates_nfe_3": 12, "loss_nfe_4": 0.33358815908432005, "updates_nfe_4": 5, "loss_nfe_2": 0.36516515566752505, "updates_nfe_2": 13, "loss_nfe_1": 0.39299650117754936, "updates_nfe_1": 8, "loss_nfe_5": 0.3845924846827984, "updates_nfe_5": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51313028107292, "grad_norm": 0.023949234522252946, "grad_rms": 6.973783166881904e-05, "grad_to_weight": 0.0014503146353603935, "update_to_weight": 2.5283568428038632e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.260111787794155, "grad_norm": 0.020004732925387972, "grad_rms": 5.8251828300988645e-05, "grad_to_weight": 0.0012302949196453103, "update_to_weight": 3.279244206850721e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067174566477674, "grad_norm": 0.002975495961823132, "grad_rms": 5.326992289764533e-05, "grad_to_weight": 0.0005872100759085309, "update_to_weight": 2.1555132044140578e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.869959761312872, "grad_norm": 0.012785333915005446, "grad_rms": 4.091312089684569e-05, "grad_to_weight": 0.0009934245446079938, "update_to_weight": 4.0736611612158316e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040914221907268, "grad_norm": 0.023822765793193684, "grad_rms": 0.0005419866884091451, "grad_to_weight": 0.004725882001654089, "update_to_weight": 1.6296101276652952e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645299006143034, "grad_norm": 0.07434311285383281, "grad_rms": 0.00017914055105092852, "grad_to_weight": 0.004466312850637054, "update_to_weight": 4.807588595778918e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.575837763091872, "grad_norm": 0.035745381710237945, "grad_rms": 4.931612644987551e-05, "grad_to_weight": 0.0014544929070096871, "update_to_weight": 6.839465948617244e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83444107402721, "grad_norm": 0.027244755881287587, "grad_rms": 3.307449893615162e-05, "grad_to_weight": 0.0008835819600516993, "update_to_weight": 6.381278816790799e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63761545338932, "grad_norm": 0.019310085728993023, "grad_rms": 1.733470716659662e-05, "grad_to_weight": 0.0005574888881995343, "update_to_weight": 7.77793545934039e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.089469990547016, "grad_norm": 0.015700679466746467, "grad_rms": 1.2629027034328312e-05, "grad_to_weight": 0.00042331905715417594, "update_to_weight": 8.275275575511724e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67316632941038, "grad_norm": 0.006164647269202019, "grad_rms": 1.113761500753832e-05, "grad_to_weight": 0.0002498522964948268, "update_to_weight": 5.400926768432358e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3652314653735, "grad_norm": 0.014990132367225023, "grad_rms": 1.2057490079636286e-05, "grad_to_weight": 0.0004011786299548669, "update_to_weight": 8.417562993637455e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51292468122138, "grad_norm": 0.02837618298863621, "grad_rms": 1.2598042066389394e-05, "grad_to_weight": 0.0004467780869966155, "update_to_weight": 8.855391848977086e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78626417552535, "grad_norm": 0.03152116458881175, "grad_rms": 2.3520179142655434e-05, "grad_to_weight": 0.0006206631871930886, "update_to_weight": 6.288290890033352e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11238028037993, "grad_norm": 0.03170425167955902, "grad_rms": 2.722636007227938e-05, "grad_to_weight": 0.000790385697830713, "update_to_weight": 6.701584105746625e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20563094822557, "grad_norm": 0.06452294382806462, "grad_rms": 0.00010058212596291095, "grad_to_weight": 0.0022875908695856985, "update_to_weight": 4.696318896662285e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417254537567182, "grad_norm": 0.003843211653203023, "grad_rms": 0.00032715590188772003, "grad_to_weight": 0.00045658731550178556, "update_to_weight": 2.7164611338909674e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.323790894652049, "grad_norm": 0.04545647299740821, "grad_rms": 0.0010523008367886836, "grad_to_weight": 0.03433810670631345, "update_to_weight": 6.890014397436045e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49300}
{"step": 49350, "lr": 1.004170517302248e-05, "window_steps": 50, "elapsed_seconds": 4818.852143466999, "steps_per_second": 1.7327777967456914, "eta_seconds": 375.1202267369521, "images_seen": 50534400, "images_per_second": 1774.364463867588, "effective_batch_size": 1024, "loss": 0.3522323668003082, "edge_loss": 0.1807836890220642, "diag_loss": 0.17144864797592163, "jvp_rms": 1.8862109184265137, "grad_norm": 0.15592190623283386, "grad_norm_postclip": 0.15592190623283386, "clip_fraction": 0.0, "ema_updates": 49350, "ema_decay": 0.9998000100000001, "loss_nfe_4": 0.31805062890052793, "updates_nfe_4": 10, "loss_nfe_2": 0.3774883236203875, "updates_nfe_2": 14, "loss_nfe_5": 0.32951241234938305, "updates_nfe_5": 12, "loss_nfe_3": 0.3438764251768589, "updates_nfe_3": 8, "loss_nfe_1": 0.4068524241447449, "updates_nfe_1": 6, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.513225456894677, "grad_norm": 0.03890667621396327, "grad_rms": 0.00011329244089540841, "grad_to_weight": 0.0023560918680316798, "update_to_weight": 2.5144145003521545e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262662209149713, "grad_norm": 0.06749931229686978, "grad_rms": 0.00019655140436101643, "grad_to_weight": 0.0041505696563563415, "update_to_weight": 4.1789602155414685e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067016094271793, "grad_norm": 0.0048126733821922535, "grad_rms": 8.616067482203843e-05, "grad_to_weight": 0.000949804242309972, "update_to_weight": 2.5279864657746708e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.870386574844499, "grad_norm": 0.03083570416023705, "grad_rms": 9.86743796160453e-05, "grad_to_weight": 0.0023958646448511674, "update_to_weight": 4.737375861324979e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040803357236634, "grad_norm": 0.06148075333632479, "grad_rms": 0.0013987355704590105, "grad_to_weight": 0.01219661807439132, "update_to_weight": 1.682968673697925e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64520733562949, "grad_norm": 0.120205672167566, "grad_rms": 0.00028965306300642545, "grad_to_weight": 0.007221638622083289, "update_to_weight": 4.8866800708450135e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.575881225060503, "grad_norm": 0.04593701719578544, "grad_rms": 6.337701936215748e-05, "grad_to_weight": 0.001869191048536749, "update_to_weight": 6.912893165546663e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.8340996174487, "grad_norm": 0.03269797573499984, "grad_rms": 3.9694580798367336e-05, "grad_to_weight": 0.0010604485339502632, "update_to_weight": 6.372000361644539e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63813527081189, "grad_norm": 0.02453850027780472, "grad_rms": 2.202826660601142e-05, "grad_to_weight": 0.0007084244023517713, "update_to_weight": 8.336035922867419e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.090217094290004, "grad_norm": 0.01989438173538237, "grad_rms": 1.6002281003158094e-05, "grad_to_weight": 0.0005363781421070487, "update_to_weight": 8.810258855493735e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67253551601583, "grad_norm": 0.0072935734821458756, "grad_rms": 1.3177236251482498e-05, "grad_to_weight": 0.00029561507682951166, "update_to_weight": 5.844799312059815e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36638631406407, "grad_norm": 0.018004672796641804, "grad_rms": 1.4482271291163645e-05, "grad_to_weight": 0.000481841424142884, "update_to_weight": 8.682082214422471e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.514095153777674, "grad_norm": 0.03043928689547185, "grad_rms": 1.3513988718412899e-05, "grad_to_weight": 0.00047925246863351383, "update_to_weight": 9.206639018044136e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78617043469487, "grad_norm": 0.03554322759810951, "grad_rms": 2.6521325951022686e-05, "grad_to_weight": 0.0006998603614701365, "update_to_weight": 6.387416033833929e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11262525786769, "grad_norm": 0.04604825465902637, "grad_rms": 3.954442371699299e-05, "grad_to_weight": 0.0011479740945151543, "update_to_weight": 6.404158928099816e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205405898072094, "grad_norm": 0.11672231613211247, "grad_rms": 0.00018195355027766715, "grad_to_weight": 0.004138295919368092, "update_to_weight": 4.630177016396346e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417360132435887, "grad_norm": 0.007858557372022512, "grad_rms": 0.0006689648285276473, "grad_to_weight": 0.0009336130625729015, "update_to_weight": 3.916126056398392e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3239614411124254, "grad_norm": 0.064623585617721, "grad_rms": 0.0014960125310577944, "grad_to_weight": 0.04881077621371106, "update_to_weight": 0.00011151635800883464, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49350}
{"step": 49400, "lr": 1.0035545678198163e-05, "window_steps": 50, "elapsed_seconds": 4846.396569396999, "steps_per_second": 1.7332465223837736, "eta_seconds": 346.1711835283571, "images_seen": 50585600, "images_per_second": 1774.8444389209842, "effective_batch_size": 1024, "loss": 0.3563637137413025, "edge_loss": 0.18396010994911194, "diag_loss": 0.17240354418754578, "jvp_rms": 2.167851448059082, "grad_norm": 0.16694892942905426, "grad_norm_postclip": 0.16694892942905426, "clip_fraction": 0.0, "ema_updates": 49400, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.3036012003819148, "updates_nfe_5": 9, "loss_nfe_2": 0.35933167636394503, "updates_nfe_2": 10, "loss_nfe_4": 0.3442973359064622, "updates_nfe_4": 11, "loss_nfe_1": 0.3896222025156021, "updates_nfe_1": 10, "loss_nfe_3": 0.38089639097452166, "updates_nfe_3": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.513301227545043, "grad_norm": 0.0061453411387936, "grad_rms": 1.7894633145224266e-05, "grad_to_weight": 0.00037214491845778555, "update_to_weight": 2.814357758071216e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26110554257603, "grad_norm": 0.004980266616029263, "grad_rms": 1.450204993448849e-05, "grad_to_weight": 0.00030626863610162054, "update_to_weight": 2.767173694126367e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066918599407097, "grad_norm": 0.0010586217708160118, "grad_rms": 1.895236990158274e-05, "grad_to_weight": 0.00020892811874654666, "update_to_weight": 2.2190617481674962e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.870202216847218, "grad_norm": 0.004268608101613901, "grad_rms": 1.3659563409416835e-05, "grad_to_weight": 0.00033166596994305586, "update_to_weight": 3.729075753678303e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040666191940241, "grad_norm": 0.010894674546445554, "grad_rms": 0.00024786242831680195, "grad_to_weight": 0.0021613560850082796, "update_to_weight": 1.4689547206440424e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645096413633073, "grad_norm": 0.030310178772607037, "grad_rms": 7.30367873948534e-05, "grad_to_weight": 0.0018209674500762672, "update_to_weight": 4.6858609671253855e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57569371716061, "grad_norm": 0.01646125543794932, "grad_rms": 2.27107759341425e-05, "grad_to_weight": 0.0006698185462188937, "update_to_weight": 5.937350518272277e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83401548988501, "grad_norm": 0.012165356067766227, "grad_rms": 1.4768458857713358e-05, "grad_to_weight": 0.00039454335980848907, "update_to_weight": 5.2144298562085435e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.6385423088393, "grad_norm": 0.008875087282545338, "grad_rms": 7.967185712175083e-06, "grad_to_weight": 0.0002562199991966906, "update_to_weight": 6.14542417304037e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.0903668423487, "grad_norm": 0.007221740808802167, "grad_rms": 5.808892545220146e-06, "grad_to_weight": 0.0001947066428190889, "update_to_weight": 6.780680354540267e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673111906929677, "grad_norm": 0.0025426575894641206, "grad_rms": 4.593797518460868e-06, "grad_to_weight": 0.00010305378579951203, "update_to_weight": 4.390132634084586e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.366846936823556, "grad_norm": 0.007402079209951612, "grad_rms": 5.9539498689027016e-06, "grad_to_weight": 0.00019809215432242303, "update_to_weight": 7.052923472342921e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51446800914241, "grad_norm": 0.017653394354495998, "grad_rms": 7.83749543698549e-06, "grad_to_weight": 0.00027794288306020183, "update_to_weight": 7.735120449029868e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78627619356803, "grad_norm": 0.022881515803264377, "grad_rms": 1.7073523702842577e-05, "grad_to_weight": 0.00045054525588868183, "update_to_weight": 5.927480881746763e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11282458627244, "grad_norm": 0.021320404104519648, "grad_rms": 1.830912159363192e-05, "grad_to_weight": 0.0005315109151355049, "update_to_weight": 6.502739843908007e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20550652167997, "grad_norm": 0.03229796021573288, "grad_rms": 5.034794307309542e-05, "grad_to_weight": 0.001145094139362726, "update_to_weight": 5.0539563138405473e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417585820031634, "grad_norm": 0.004118144149412127, "grad_rms": 0.00035055971017930053, "grad_to_weight": 0.0004892310262655155, "update_to_weight": 2.1120037583264877e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3241388527501978, "grad_norm": 0.05100673809167384, "grad_rms": 0.001180787457460458, "grad_to_weight": 0.03852068684921851, "update_to_weight": 6.831808176519086e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49400}
{"step": 49450, "lr": 1.0029877907442364e-05, "window_steps": 50, "elapsed_seconds": 4874.149777133, "steps_per_second": 1.7336356875291457, "eta_seconds": 317.2523523577692, "images_seen": 50636800, "images_per_second": 1775.2429440298451, "effective_batch_size": 1024, "loss": 0.3511733412742615, "edge_loss": 0.17958179116249084, "diag_loss": 0.17159156501293182, "jvp_rms": 1.813320279121399, "grad_norm": 0.22733347117900848, "grad_norm_postclip": 0.21504171192646027, "clip_fraction": 0.019999999552965164, "ema_updates": 49450, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.36483187334878103, "updates_nfe_2": 7, "loss_nfe_1": 0.39537911613782245, "updates_nfe_1": 6, "loss_nfe_4": 0.37591610505030704, "updates_nfe_4": 13, "loss_nfe_5": 0.3030197811978204, "updates_nfe_5": 14, "loss_nfe_3": 0.35033833235502243, "updates_nfe_3": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.512503770225724, "grad_norm": 0.0383095731983438, "grad_rms": 0.00011155373523642279, "grad_to_weight": 0.0023200341832727476, "update_to_weight": 5.652276481035481e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.261127124867258, "grad_norm": 0.012031236473801986, "grad_rms": 3.5033785451395285e-05, "grad_to_weight": 0.0007398771549730565, "update_to_weight": 5.689330784891934e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067162333051614, "grad_norm": 0.002791990160764133, "grad_rms": 4.9984642057373584e-05, "grad_to_weight": 0.0005509967862195376, "update_to_weight": 2.5150503910739678e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.869688845936055, "grad_norm": 0.010034665025798148, "grad_rms": 3.2110969184620934e-05, "grad_to_weight": 0.0007797131030846064, "update_to_weight": 5.524358586074582e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040641785586009, "grad_norm": 0.009454884738467239, "grad_rms": 0.0002151060759769592, "grad_to_weight": 0.00187573034162118, "update_to_weight": 1.9886006681702997e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64511108104691, "grad_norm": 0.027938082323709944, "grad_rms": 6.732087574293139e-05, "grad_to_weight": 0.0016784557452140928, "update_to_weight": 5.6097363491267286e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57582410360014, "grad_norm": 0.02149058314598331, "grad_rms": 2.9649489394187288e-05, "grad_to_weight": 0.0008744603255373694, "update_to_weight": 6.404665979003794e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833812593051764, "grad_norm": 0.014139701193329003, "grad_rms": 1.7165267845085208e-05, "grad_to_weight": 0.00045857777563697424, "update_to_weight": 5.192849096755219e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63870089378894, "grad_norm": 0.010084434153232538, "grad_rms": 9.052819103989953e-06, "grad_to_weight": 0.0002911319966691007, "update_to_weight": 5.943492568831256e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09087203188862, "grad_norm": 0.008314160475801516, "grad_rms": 6.687593211429314e-06, "grad_to_weight": 0.00022415651130158043, "update_to_weight": 6.291414494665044e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67338401813278, "grad_norm": 0.0033505651635491467, "grad_rms": 6.053437158637398e-06, "grad_to_weight": 0.00013579674199075306, "update_to_weight": 4.2946717529811546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.367469258555126, "grad_norm": 0.008676361422830632, "grad_rms": 6.9789338226161185e-06, "grad_to_weight": 0.00023219023377785253, "update_to_weight": 6.712335488640039e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.515438575987176, "grad_norm": 0.02105773150922453, "grad_rms": 9.348903179902067e-06, "grad_to_weight": 0.0003315372133348643, "update_to_weight": 7.623095503743994e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78654058978742, "grad_norm": 0.02880226587045259, "grad_rms": 2.149141574635497e-05, "grad_to_weight": 0.0005671239965544018, "update_to_weight": 5.93878904049393e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11304673737321, "grad_norm": 0.02581220366226177, "grad_rms": 2.216650177619087e-05, "grad_to_weight": 0.0006434864903496002, "update_to_weight": 6.556915149647363e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205376684699342, "grad_norm": 0.0365453301950722, "grad_rms": 5.6968990981439263e-05, "grad_to_weight": 0.0012956866559026337, "update_to_weight": 5.573757434139169e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417879476759804, "grad_norm": 0.003818760019718784, "grad_rms": 0.00032507444061860227, "grad_to_weight": 0.00045364869267392914, "update_to_weight": 5.877078101815628e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3244256502180964, "grad_norm": 0.05060612071525353, "grad_rms": 0.0011715133107297335, "grad_to_weight": 0.03820986154029869, "update_to_weight": 0.00010619514668365065, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49450}
{"step": 49500, "lr": 1.0024701922736924e-05, "window_steps": 50, "elapsed_seconds": 4901.631927126, "steps_per_second": 1.7341163364308039, "eta_seconds": 288.33128983094116, "images_seen": 50688000, "images_per_second": 1775.7351285051432, "effective_batch_size": 1024, "loss": 0.35982969403266907, "edge_loss": 0.18546539545059204, "diag_loss": 0.17436428368091583, "jvp_rms": 1.8934375047683716, "grad_norm": 0.3016242980957031, "grad_norm_postclip": 0.2804275155067444, "clip_fraction": 0.03999999910593033, "ema_updates": 49500, "ema_decay": 0.9998000100000001, "loss_nfe_5": 0.36467018723487854, "updates_nfe_5": 8, "loss_nfe_2": 0.3452421592341529, "updates_nfe_2": 9, "loss_nfe_1": 0.3957562930881977, "updates_nfe_1": 8, "loss_nfe_4": 0.32573975508029646, "updates_nfe_4": 13, "loss_nfe_3": 0.3805230011542638, "updates_nfe_3": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.514173489018173, "grad_norm": 0.07580949094882902, "grad_rms": 0.00022074983289240268, "grad_to_weight": 0.0045905713052634385, "update_to_weight": 3.244492904637884e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262252179426614, "grad_norm": 0.030578502875640803, "grad_rms": 8.90416135949773e-05, "grad_to_weight": 0.0018803362866507314, "update_to_weight": 4.379415881287314e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067145959035127, "grad_norm": 0.001948414919313133, "grad_rms": 3.48822226130119e-05, "grad_to_weight": 0.00038451920174885693, "update_to_weight": 2.0457520589458976e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.870917108044338, "grad_norm": 0.022632950896253782, "grad_rms": 7.242553557275696e-05, "grad_to_weight": 0.0017584567367082305, "update_to_weight": 4.92079677870793e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040443692431961, "grad_norm": 0.034472289367106695, "grad_rms": 0.0007842717389807838, "grad_to_weight": 0.006839137875672642, "update_to_weight": 2.0690225410744554e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645507096335052, "grad_norm": 0.12500654143860243, "grad_rms": 0.0003012214562808348, "grad_to_weight": 0.0075099268958964864, "update_to_weight": 6.296170443201222e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.576250027815217, "grad_norm": 0.02521501619012612, "grad_rms": 3.4787904545211815e-05, "grad_to_weight": 0.0010259911972570247, "update_to_weight": 7.249633486196256e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83378883918886, "grad_norm": 0.008657735512038458, "grad_rms": 1.0510289217862693e-05, "grad_to_weight": 0.0002807872738959257, "update_to_weight": 6.650441102024997e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63868679737838, "grad_norm": 0.004884801399953341, "grad_rms": 4.38509714682598e-06, "grad_to_weight": 0.0001410215528240824, "update_to_weight": 8.103103150232857e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09088025967434, "grad_norm": 0.0040676346082130525, "grad_rms": 3.271849956665421e-06, "grad_to_weight": 0.0001096667045843998, "update_to_weight": 8.526884039563105e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67329743762095, "grad_norm": 0.0013932933713779441, "grad_rms": 2.517251107048558e-06, "grad_to_weight": 5.646968650625112e-05, "update_to_weight": 5.6560718980537874e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36776816530859, "grad_norm": 0.0042089225060303595, "grad_rms": 3.3854965466067886e-06, "grad_to_weight": 0.0001126351054044975, "update_to_weight": 8.207906159199326e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51610547173744, "grad_norm": 0.011512636268752952, "grad_rms": 5.111211612459387e-06, "grad_to_weight": 0.00018125538685421594, "update_to_weight": 8.21882489636453e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78667519097037, "grad_norm": 0.02459752293181521, "grad_rms": 1.8353958471040042e-05, "grad_to_weight": 0.0004843302468476718, "update_to_weight": 5.905547902684847e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11358993727998, "grad_norm": 0.04401480954487692, "grad_rms": 3.779818130683874e-05, "grad_to_weight": 0.0010972543124087555, "update_to_weight": 7.436143792852691e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.2057077678188, "grad_norm": 0.0734655344493447, "grad_rms": 0.00011452235749824136, "grad_to_weight": 0.0026046336101221658, "update_to_weight": 5.828051027515451e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417741713150647, "grad_norm": 0.004637780856839433, "grad_rms": 0.00039479412426124725, "grad_to_weight": 0.0005509530958397123, "update_to_weight": 3.051654521822391e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.32432276674522, "grad_norm": 0.05107732119719834, "grad_rms": 0.0011824214307124261, "grad_to_weight": 0.03856863483720873, "update_to_weight": 6.550125888330484e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49500}
{"step": 49550, "lr": 1.0020017780685572e-05, "window_steps": 50, "elapsed_seconds": 4931.2430234309995, "steps_per_second": 1.7338427571657553, "eta_seconds": 259.5391064963684, "images_seen": 50739200, "images_per_second": 1775.4549833377334, "effective_batch_size": 1024, "loss": 0.34218236804008484, "edge_loss": 0.17588864266872406, "diag_loss": 0.16629381477832794, "jvp_rms": 1.6442577838897705, "grad_norm": 0.16753114759922028, "grad_norm_postclip": 0.16753114759922028, "clip_fraction": 0.0, "ema_updates": 49550, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.3251148536801338, "updates_nfe_3": 18, "loss_nfe_5": 0.3335349178976483, "updates_nfe_5": 9, "loss_nfe_2": 0.3400717079639435, "updates_nfe_2": 7, "loss_nfe_1": 0.39378273487091064, "updates_nfe_1": 8, "loss_nfe_4": 0.34055955708026886, "updates_nfe_4": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.5136440391915, "grad_norm": 0.09550830903343517, "grad_rms": 0.00027811086705750303, "grad_to_weight": 0.005783599840638881, "update_to_weight": 4.4477505192992895e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.261010767827635, "grad_norm": 0.03548919949212248, "grad_rms": 0.00010334108248608666, "grad_to_weight": 0.002182471926181721, "update_to_weight": 5.711274366490479e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066915964384719, "grad_norm": 0.003163965713178158, "grad_rms": 5.664407270394153e-05, "grad_to_weight": 0.0006244361926303157, "update_to_weight": 1.9683207648314648e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.870548997387043, "grad_norm": 0.020300173773447172, "grad_rms": 6.496063922470237e-05, "grad_to_weight": 0.001577257798215793, "update_to_weight": 4.9334423601497404e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040294786339697, "grad_norm": 0.026974030521140564, "grad_rms": 0.0006136804434091815, "grad_to_weight": 0.005351677166630432, "update_to_weight": 1.664500336958328e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645107414194662, "grad_norm": 0.05309048260564096, "grad_rms": 0.00012792924515056715, "grad_to_weight": 0.0031895548213985276, "update_to_weight": 4.322527607215022e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.576030236940856, "grad_norm": 0.02321681167062612, "grad_rms": 3.203108108882247e-05, "grad_to_weight": 0.0009446933229976394, "update_to_weight": 5.829829659089642e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83411743266803, "grad_norm": 0.014055246769094844, "grad_rms": 1.706274214154589e-05, "grad_to_weight": 0.0004558342491815134, "update_to_weight": 5.389936288410663e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63912378343851, "grad_norm": 0.011004344089413776, "grad_rms": 9.878624312063253e-06, "grad_to_weight": 0.00031768540561857747, "update_to_weight": 6.664787891841698e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09145784567098, "grad_norm": 0.008832390516793888, "grad_rms": 7.1044376678463295e-06, "grad_to_weight": 0.00023812465267726687, "update_to_weight": 6.991749520348626e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673407518504995, "grad_norm": 0.003481396549832685, "grad_rms": 6.289809094888862e-06, "grad_to_weight": 0.00014109913870719503, "update_to_weight": 4.5412301114040025e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.368432938423275, "grad_norm": 0.0087949555892709, "grad_rms": 7.074326441595438e-06, "grad_to_weight": 0.0002353578916130486, "update_to_weight": 7.053396600530965e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51696646799873, "grad_norm": 0.020685216939729925, "grad_rms": 9.183519617965096e-06, "grad_to_weight": 0.00032566443408709703, "update_to_weight": 7.959682903106923e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.787076588808496, "grad_norm": 0.029100519039981326, "grad_rms": 2.171396361438871e-05, "grad_to_weight": 0.0005729906305808898, "update_to_weight": 5.8769038083415164e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.113568635461434, "grad_norm": 0.03237300530455439, "grad_rms": 2.780065929175931e-05, "grad_to_weight": 0.0008070337894578598, "update_to_weight": 6.098981356923067e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20531501192417, "grad_norm": 0.04936435464741105, "grad_rms": 7.695203353483815e-05, "grad_to_weight": 0.0017501791639817394, "update_to_weight": 4.618695252364373e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417853193087362, "grad_norm": 0.003863405397926198, "grad_rms": 0.00032887490759532783, "grad_to_weight": 0.0004589537628309768, "update_to_weight": 1.618803194002038e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3245494058365253, "grad_norm": 0.04585359289964915, "grad_rms": 0.001061494018262693, "grad_to_weight": 0.03461825787516783, "update_to_weight": 5.097626626199124e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49550}
{"step": 49600, "lr": 1.0015825532513342e-05, "window_steps": 50, "elapsed_seconds": 4958.714778838999, "steps_per_second": 1.7343203599247035, "eta_seconds": 230.63789669018604, "images_seen": 50790400, "images_per_second": 1775.9440485628963, "effective_batch_size": 1024, "loss": 0.3498483896255493, "edge_loss": 0.18211467564105988, "diag_loss": 0.16773366928100586, "jvp_rms": 1.7157421112060547, "grad_norm": 0.18225105106830597, "grad_norm_postclip": 0.18225105106830597, "clip_fraction": 0.0, "ema_updates": 49600, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39789979159832, "updates_nfe_1": 12, "loss_nfe_4": 0.31108903655639064, "updates_nfe_4": 13, "loss_nfe_3": 0.3350561906894048, "updates_nfe_3": 12, "loss_nfe_5": 0.387542466322581, "updates_nfe_5": 6, "loss_nfe_2": 0.332505145243236, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51269505260697, "grad_norm": 0.02777686319101803, "grad_rms": 8.088351248567611e-05, "grad_to_weight": 0.0016821520110754248, "update_to_weight": 3.352815543234214e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262807640549514, "grad_norm": 0.008875601293638994, "grad_rms": 2.5844884035864442e-05, "grad_to_weight": 0.0005457607007235738, "update_to_weight": 3.835065962350309e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067086296941213, "grad_norm": 0.001103816020369626, "grad_rms": 1.976147675974257e-05, "grad_to_weight": 0.0002178403831479944, "update_to_weight": 1.7525108954010276e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.871185626568897, "grad_norm": 0.00921112350475449, "grad_rms": 2.947563294404868e-05, "grad_to_weight": 0.0007156390850071145, "update_to_weight": 4.5279448284145486e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0398461493972935, "grad_norm": 0.014745537482121065, "grad_rms": 0.0003354725936579141, "grad_to_weight": 0.002925791193821355, "update_to_weight": 2.1046508963366232e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645269671633628, "grad_norm": 0.04459095171530612, "grad_rms": 0.00010744838836477983, "grad_to_weight": 0.002678896322797142, "update_to_weight": 5.655100053360328e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57633570847104, "grad_norm": 0.021112605283910426, "grad_rms": 2.9128012124974046e-05, "grad_to_weight": 0.0008590623734291388, "update_to_weight": 6.564942327127044e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833796757145194, "grad_norm": 0.011746160913915399, "grad_rms": 1.4259565706661056e-05, "grad_to_weight": 0.0003809508445045267, "update_to_weight": 5.890452458751033e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63945856741629, "grad_norm": 0.00820155254616931, "grad_rms": 7.362552072248818e-06, "grad_to_weight": 0.0002367690744994532, "update_to_weight": 6.778388301168088e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09171783871114, "grad_norm": 0.007043084921089668, "grad_rms": 5.66518857109968e-06, "grad_to_weight": 0.0001898829531626352, "update_to_weight": 7.240658049051127e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673017904180266, "grad_norm": 0.0025319265543728136, "grad_rms": 4.574409849992558e-06, "grad_to_weight": 0.00010261924845212542, "update_to_weight": 4.98124237550171e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36908953318321, "grad_norm": 0.007338845033415492, "grad_rms": 5.903086711887138e-06, "grad_to_weight": 0.00019638811448426392, "update_to_weight": 7.217215332536019e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.517310478866705, "grad_norm": 0.01744008000223687, "grad_rms": 7.742791255517487e-06, "grad_to_weight": 0.0002745720791820914, "update_to_weight": 7.831667712066047e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78692276007181, "grad_norm": 0.02544595805451092, "grad_rms": 1.8987036161443875e-05, "grad_to_weight": 0.0005010336651961179, "update_to_weight": 6.0531869955742156e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.114044880561316, "grad_norm": 0.026776477210522346, "grad_rms": 2.299458184250083e-05, "grad_to_weight": 0.0006675087812821847, "update_to_weight": 7.405029070435602e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205569276006358, "grad_norm": 0.05411405574035338, "grad_rms": 8.435614446458013e-05, "grad_to_weight": 0.0019185592466090234, "update_to_weight": 6.194639179604346e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417729930632046, "grad_norm": 0.006958976581960175, "grad_rms": 0.0005923874262790793, "grad_to_weight": 0.0008267046625761324, "update_to_weight": 4.699178697202711e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3243541817066238, "grad_norm": 0.09330475897531934, "grad_rms": 0.002159971275195354, "grad_to_weight": 0.0704530255305892, "update_to_weight": 0.0001503198962094345, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49600}
{"step": 49650, "lr": 1.0012125224065987e-05, "window_steps": 50, "elapsed_seconds": 4986.205454708, "steps_per_second": 1.7347861171329448, "eta_seconds": 201.75397793616182, "images_seen": 50841600, "images_per_second": 1776.4209839441355, "effective_batch_size": 1024, "loss": 0.35978466272354126, "edge_loss": 0.18586324155330658, "diag_loss": 0.17392133176326752, "jvp_rms": 1.904140591621399, "grad_norm": 0.19297660887241364, "grad_norm_postclip": 0.19297660887241364, "clip_fraction": 0.0, "ema_updates": 49650, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.3936563700437546, "updates_nfe_1": 10, "loss_nfe_5": 0.3889050292117255, "updates_nfe_5": 14, "loss_nfe_3": 0.3118117054303487, "updates_nfe_3": 6, "loss_nfe_4": 0.3025529781977336, "updates_nfe_4": 9, "loss_nfe_2": 0.3649227564985102, "updates_nfe_2": 11, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51485629715333, "grad_norm": 0.013411148621901141, "grad_rms": 3.905195484267504e-05, "grad_to_weight": 0.000812065717109075, "update_to_weight": 3.44106748562093e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26055283848865, "grad_norm": 0.008274367676069932, "grad_rms": 2.4094150467461353e-05, "grad_to_weight": 0.0005088613996250204, "update_to_weight": 7.043508505832101e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.066977510621065, "grad_norm": 0.0006255315021873175, "grad_rms": 1.1198810322414202e-05, "grad_to_weight": 0.00012345259099258275, "update_to_weight": 2.1853947301454902e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.871180884568489, "grad_norm": 0.008788746258556393, "grad_rms": 2.812402402615425e-05, "grad_to_weight": 0.0006828236148163681, "update_to_weight": 6.071912122376194e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.040020235111851, "grad_norm": 0.021734384885241688, "grad_rms": 0.0004944743776109938, "grad_to_weight": 0.0043123606397106745, "update_to_weight": 2.098204609549916e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64562626575813, "grad_norm": 0.09449208626617522, "grad_rms": 0.00022769243516821344, "grad_to_weight": 0.005676691567955948, "update_to_weight": 4.8450860896071985e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.576896971466795, "grad_norm": 0.020128247862820883, "grad_rms": 2.7769943117798736e-05, "grad_to_weight": 0.0008189906108240316, "update_to_weight": 6.304365713242475e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833864059691944, "grad_norm": 0.007961714146991228, "grad_rms": 9.665335495461952e-06, "grad_to_weight": 0.000258213311558291, "update_to_weight": 5.2438065402773565e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63965238824034, "grad_norm": 0.0040531079909304655, "grad_rms": 3.6384841125734476e-06, "grad_to_weight": 0.00011700775589498804, "update_to_weight": 6.223960685919964e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09191530056196, "grad_norm": 0.003403779869841817, "grad_rms": 2.737870554339937e-06, "grad_to_weight": 9.176608547335512e-05, "update_to_weight": 6.375375241187296e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67292019041638, "grad_norm": 0.001177194963523202, "grad_rms": 2.12682797895616e-06, "grad_to_weight": 4.771202413164113e-05, "update_to_weight": 4.3292033415911366e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.369672618858, "grad_norm": 0.0035369488057402027, "grad_rms": 2.844986561335394e-06, "grad_to_weight": 9.46475726938892e-05, "update_to_weight": 6.657714058399427e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.518040775092004, "grad_norm": 0.009596905721712222, "grad_rms": 4.260693625979268e-06, "grad_to_weight": 0.00015108944804663368, "update_to_weight": 7.617134478723187e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78707899237881, "grad_norm": 0.024062995450768128, "grad_rms": 1.7955109561905447e-05, "grad_to_weight": 0.00047380152448576654, "update_to_weight": 5.958962376498585e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11433092921328, "grad_norm": 0.037958394665034756, "grad_rms": 3.2597171236255454e-05, "grad_to_weight": 0.0009462552106880969, "update_to_weight": 6.701271510145658e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205621210515712, "grad_norm": 0.08929139325695186, "grad_rms": 0.00013919262871679703, "grad_to_weight": 0.003165730426233688, "update_to_weight": 5.466067944763357e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.417876304596795, "grad_norm": 0.007283690947361749, "grad_rms": 0.0006200289486970066, "grad_to_weight": 0.0008652646681662813, "update_to_weight": 4.025357175058145e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3245006250025022, "grad_norm": 0.11569207427179773, "grad_rms": 0.002678229491605578, "grad_to_weight": 0.08734769322708259, "update_to_weight": 0.00014041231298360488, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49650}
{"step": 49700, "lr": 1.0008916895809486e-05, "window_steps": 50, "elapsed_seconds": 5013.763970245999, "steps_per_second": 1.7352232876596974, "eta_seconds": 172.88841276710343, "images_seen": 50892800, "images_per_second": 1776.86864656353, "effective_batch_size": 1024, "loss": 0.3890533447265625, "edge_loss": 0.2008022516965866, "diag_loss": 0.1882510781288147, "jvp_rms": 2.8324999809265137, "grad_norm": 0.19757461547851562, "grad_norm_postclip": 0.19757461547851562, "clip_fraction": 0.0, "ema_updates": 49700, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.39778060615062716, "updates_nfe_1": 10, "loss_nfe_4": 0.41933912309733307, "updates_nfe_4": 11, "loss_nfe_2": 0.3742019661835262, "updates_nfe_2": 14, "loss_nfe_3": 0.387847475707531, "updates_nfe_3": 10, "loss_nfe_5": 0.348965585231781, "updates_nfe_5": 5, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51205466375248, "grad_norm": 0.3162483009473672, "grad_rms": 0.0009208841625616525, "grad_to_weight": 0.01915257109956161, "update_to_weight": 6.580699042795605e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.263964477560314, "grad_norm": 0.27351921767812026, "grad_rms": 0.0007964612456778197, "grad_to_weight": 0.016817499697290884, "update_to_weight": 6.046247595011765e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067062770685825, "grad_norm": 0.010023244281691013, "grad_rms": 0.00017944485790624066, "grad_to_weight": 0.0019781172516113057, "update_to_weight": 3.076513597063303e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.87177006474282, "grad_norm": 0.11484023757026744, "grad_rms": 0.00036748923061137204, "grad_to_weight": 0.00892186832056823, "update_to_weight": 7.603454931223658e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.039637806306124, "grad_norm": 0.03475264497144506, "grad_rms": 0.0007906500498323184, "grad_to_weight": 0.006895861628777946, "update_to_weight": 1.898146114291058e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645090913349552, "grad_norm": 0.23062650486214778, "grad_rms": 0.0005557281311206865, "grad_to_weight": 0.013855526897554083, "update_to_weight": 6.91435092222909e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57642759753808, "grad_norm": 0.028117267966387166, "grad_rms": 3.879199706680522e-05, "grad_to_weight": 0.0011440746566927322, "update_to_weight": 8.006834533247142e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833990746440243, "grad_norm": 0.012555645346117634, "grad_rms": 1.5242260949311138e-05, "grad_to_weight": 0.0004072014371855895, "update_to_weight": 6.451865078854826e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63952199980537, "grad_norm": 0.004061949807931735, "grad_rms": 3.6464214314796767e-06, "grad_to_weight": 0.00011726344861094095, "update_to_weight": 7.19085771825222e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09268374625918, "grad_norm": 0.00302011957962699, "grad_rms": 2.4292688669172114e-06, "grad_to_weight": 8.142089691559648e-05, "update_to_weight": 7.250768475009507e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673334543591793, "grad_norm": 0.0012819899797239494, "grad_rms": 2.3161602301270773e-06, "grad_to_weight": 5.195852135263616e-05, "update_to_weight": 5.192502014621805e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.3703324583402, "grad_norm": 0.0030343728792831853, "grad_rms": 2.4407336768999716e-06, "grad_to_weight": 8.119737448592013e-05, "update_to_weight": 6.803302733486663e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51883256045997, "grad_norm": 0.008560564420125294, "grad_rms": 3.800593995322201e-06, "grad_to_weight": 0.0001347720679843569, "update_to_weight": 7.33299319257811e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78720878500683, "grad_norm": 0.03059068506245957, "grad_rms": 2.2825882296905507e-05, "grad_to_weight": 0.0006023305039652112, "update_to_weight": 6.0471849720600446e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11461545430637, "grad_norm": 0.12915720498004016, "grad_rms": 0.00011091511019586633, "grad_to_weight": 0.003219704427359154, "update_to_weight": 7.80444677512318e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205375602721993, "grad_norm": 0.31024863358995164, "grad_rms": 0.00048363365482392216, "grad_to_weight": 0.010999627800028684, "update_to_weight": 6.418377175865497e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.418083852171318, "grad_norm": 0.002868392037187744, "grad_rms": 0.00024417374544324585, "grad_to_weight": 0.000340741680358516, "update_to_weight": 5.510930410380722e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.324792338238781, "grad_norm": 0.04624686490690383, "grad_rms": 0.001070598122365606, "grad_to_weight": 0.03490876537554996, "update_to_weight": 9.126439059534455e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49700}
{"step": 49750, "lr": 1.0006200582829606e-05, "window_steps": 50, "elapsed_seconds": 5041.327849722, "steps_per_second": 1.7356538318535486, "eta_seconds": 144.0379385634857, "images_seen": 50944000, "images_per_second": 1777.3095238180338, "effective_batch_size": 1024, "loss": 0.3694321811199188, "edge_loss": 0.19336764514446259, "diag_loss": 0.17606458067893982, "jvp_rms": 2.0672264099121094, "grad_norm": 0.26698750257492065, "grad_norm_postclip": 0.2641989588737488, "clip_fraction": 0.019999999552965164, "ema_updates": 49750, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.344769428173701, "updates_nfe_2": 3, "loss_nfe_5": 0.34226437471807003, "updates_nfe_5": 16, "loss_nfe_4": 0.3800474920056083, "updates_nfe_4": 11, "loss_nfe_3": 0.3459723014384508, "updates_nfe_3": 8, "loss_nfe_1": 0.4177309622367223, "updates_nfe_1": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.516380733064107, "grad_norm": 0.19490628915981764, "grad_rms": 0.00056754807641104, "grad_to_weight": 0.011800786886053986, "update_to_weight": 8.727044955605102e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.2617764569298, "grad_norm": 0.06174957065671917, "grad_rms": 0.00017980871831535392, "grad_to_weight": 0.0037972217131545423, "update_to_weight": 0.000111560459763895, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.06710982308737, "grad_norm": 0.00598109679763798, "grad_rms": 0.00010707880949645444, "grad_to_weight": 0.00118037638939385, "update_to_weight": 4.100642297163651e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.872293428536794, "grad_norm": 0.0369914496337147, "grad_rms": 0.00011837279034515565, "grad_to_weight": 0.0028737264139510493, "update_to_weight": 9.088966772401274e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.039883237727513, "grad_norm": 0.06352137670949742, "grad_rms": 0.0014451613597194718, "grad_to_weight": 0.012603739752141413, "update_to_weight": 3.20981027455519e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.6462074336417, "grad_norm": 0.10831279206470441, "grad_rms": 0.00026099543739155443, "grad_to_weight": 0.006506754916786999, "update_to_weight": 7.596663239034563e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.577463188403797, "grad_norm": 0.03422343995993365, "grad_rms": 4.72163790638837e-05, "grad_to_weight": 0.0013924724328783062, "update_to_weight": 0.0001015785525351448, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833723515971517, "grad_norm": 0.017219053302865488, "grad_rms": 2.0903529568356973e-05, "grad_to_weight": 0.0005584487158661267, "update_to_weight": 9.11459279886374e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63795201403458, "grad_norm": 0.010898236251283817, "grad_rms": 9.783371068350213e-06, "grad_to_weight": 0.0003146328122074907, "update_to_weight": 0.00011210621215966112, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.092142380390854, "grad_norm": 0.008435933261658262, "grad_rms": 6.785542590492015e-06, "grad_to_weight": 0.0002274318149419702, "update_to_weight": 0.00010885424052614315, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673704363384893, "grad_norm": 0.0035751831919946047, "grad_rms": 6.459252611708995e-06, "grad_to_weight": 0.00014489851784477403, "update_to_weight": 7.0557918275794e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.36988657798633, "grad_norm": 0.008661297653413285, "grad_rms": 6.966817101705186e-06, "grad_to_weight": 0.0002317721151049851, "update_to_weight": 0.00010598858270191962, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.518161849707404, "grad_norm": 0.02275433248317146, "grad_rms": 1.0102135228349888e-05, "grad_to_weight": 0.0003582334850465494, "update_to_weight": 0.00011162119750147598, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78745635150721, "grad_norm": 0.031771700943633005, "grad_rms": 2.370712210697838e-05, "grad_to_weight": 0.0006255816539370774, "update_to_weight": 8.291183296605918e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11499735379915, "grad_norm": 0.04141108286741808, "grad_rms": 3.55622035973885e-05, "grad_to_weight": 0.0010323092508816077, "update_to_weight": 9.579218325497587e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205836521315717, "grad_norm": 0.0970594288732526, "grad_rms": 0.00015130189544407393, "grad_to_weight": 0.0034411115160474974, "update_to_weight": 8.293357253689042e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.418100165737554, "grad_norm": 0.013824209206718757, "grad_rms": 0.0011767948369794561, "grad_to_weight": 0.0016422006075651804, "update_to_weight": 7.017479030166206e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3246100193698207, "grad_norm": 0.15164695235213912, "grad_rms": 0.0035105718577353832, "grad_to_weight": 0.11448422564724726, "update_to_weight": 0.00016167327436523586, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49750}
{"step": 49800, "lr": 1.0003976314831507e-05, "window_steps": 50, "elapsed_seconds": 5068.914833641999, "steps_per_second": 1.736071780412461, "eta_seconds": 115.20260985549999, "images_seen": 50995200, "images_per_second": 1777.73750314236, "effective_batch_size": 1024, "loss": 0.36194413900375366, "edge_loss": 0.1867786943912506, "diag_loss": 0.17516544461250305, "jvp_rms": 2.163515567779541, "grad_norm": 0.2328273355960846, "grad_norm_postclip": 0.2328273355960846, "clip_fraction": 0.0, "ema_updates": 49800, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.41563307245572406, "updates_nfe_1": 6, "loss_nfe_3": 0.4324966503514184, "updates_nfe_3": 9, "loss_nfe_4": 0.32557201843995315, "updates_nfe_4": 13, "loss_nfe_2": 0.3646954968571663, "updates_nfe_2": 12, "loss_nfe_5": 0.3102156803011894, "updates_nfe_5": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.511975190964623, "grad_norm": 0.0032699229723112864, "grad_rms": 9.521696303118295e-06, "grad_to_weight": 0.00019803342328788098, "update_to_weight": 3.1811366103413804e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26526570358453, "grad_norm": 0.004031932224815274, "grad_rms": 1.1740592816567556e-05, "grad_to_weight": 0.0002478860350819058, "update_to_weight": 3.356274331675513e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067325129302881, "grad_norm": 0.0006008211109043886, "grad_rms": 1.075642335388821e-05, "grad_to_weight": 0.00011856770496726435, "update_to_weight": 1.8773923408846008e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.872403077371688, "grad_norm": 0.006156939703856192, "grad_rms": 1.970223227121326e-05, "grad_to_weight": 0.00047830538453845, "update_to_weight": 4.3378323129969996e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.0395248319277295, "grad_norm": 0.012987620742036268, "grad_rms": 0.00029547860300508353, "grad_to_weight": 0.0025771518496652025, "update_to_weight": 1.5692270525825718e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645350341410097, "grad_norm": 0.06462476734962497, "grad_rms": 0.0001557227830547186, "grad_to_weight": 0.0038824516170652334, "update_to_weight": 4.4962573859200096e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.576872137091993, "grad_norm": 0.016748522940784393, "grad_rms": 2.310710462940153e-05, "grad_to_weight": 0.0006814749593585235, "update_to_weight": 5.650122936911529e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.834293604838347, "grad_norm": 0.008382054014967895, "grad_rms": 1.0175618302795345e-05, "grad_to_weight": 0.00027184193425636415, "update_to_weight": 4.624946282191539e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.640085838158654, "grad_norm": 0.005182065106905585, "grad_rms": 4.651951441705533e-06, "grad_to_weight": 0.00014959735178245867, "update_to_weight": 4.9023519897457816e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09364962865527, "grad_norm": 0.004386995770651593, "grad_rms": 3.52873188096004e-06, "grad_to_weight": 0.00011826810827647944, "update_to_weight": 5.302710735112421e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673701889692317, "grad_norm": 0.0016012198872476608, "grad_rms": 2.892910148432037e-06, "grad_to_weight": 6.489581070591544e-05, "update_to_weight": 3.511090002545583e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.371044549264674, "grad_norm": 0.0045088081765730205, "grad_rms": 3.626713128889883e-06, "grad_to_weight": 0.00012064977661058547, "update_to_weight": 5.6127694266594226e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.51975885994058, "grad_norm": 0.010952084207792886, "grad_rms": 4.8623459194518305e-06, "grad_to_weight": 0.00017242011626558514, "update_to_weight": 6.475414707739125e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78774958405324, "grad_norm": 0.02048273483310821, "grad_rms": 1.5283622889276456e-05, "grad_to_weight": 0.0004033006975276485, "update_to_weight": 5.284544842435745e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.115224057356016, "grad_norm": 0.027613981956585067, "grad_rms": 2.371379786466129e-05, "grad_to_weight": 0.0006883666389873106, "update_to_weight": 6.175702714798897e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20543619338965, "grad_norm": 0.06015271472141468, "grad_rms": 9.376955808530193e-05, "grad_to_weight": 0.0021326638705028195, "update_to_weight": 5.1713392895333575e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.418096993657706, "grad_norm": 0.00656633404153656, "grad_rms": 0.0005589634736001485, "grad_to_weight": 0.0007800259425002722, "update_to_weight": 2.387601391656848e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3248201878015027, "grad_norm": 0.0898675783320135, "grad_rms": 0.0020804017919371426, "grad_to_weight": 0.0678337929626102, "update_to_weight": 5.30316612127819e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49800}
{"step": 49850, "lr": 1.0002244116139426e-05, "window_steps": 50, "elapsed_seconds": 5096.423439079, "steps_per_second": 1.7365119099286082, "eta_seconds": 86.38005828947458, "images_seen": 51046400, "images_per_second": 1778.1881957668948, "effective_batch_size": 1024, "loss": 0.38234296441078186, "edge_loss": 0.1963135153055191, "diag_loss": 0.18602944910526276, "jvp_rms": 2.604804515838623, "grad_norm": 0.1593082994222641, "grad_norm_postclip": 0.1593082994222641, "clip_fraction": 0.0, "ema_updates": 49850, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.34136271689619335, "updates_nfe_3": 14, "loss_nfe_4": 0.3732438501384523, "updates_nfe_4": 9, "loss_nfe_1": 0.39369241065449184, "updates_nfe_1": 9, "loss_nfe_2": 0.36934906740983325, "updates_nfe_2": 6, "loss_nfe_5": 0.43496256445844966, "updates_nfe_5": 12, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.513755844146363, "grad_norm": 0.04673182705951737, "grad_rms": 0.00013607851583000857, "grad_to_weight": 0.0028298727134252995, "update_to_weight": 4.2848412679557655e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262075779019668, "grad_norm": 0.0221123256458341, "grad_rms": 6.43889324421808e-05, "grad_to_weight": 0.001359748038707461, "update_to_weight": 4.924991956032244e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067256247365515, "grad_norm": 0.0039811177458735775, "grad_rms": 7.12734408280602e-05, "grad_to_weight": 0.0007856555010304393, "update_to_weight": 2.2414923222176946e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.87140020025927, "grad_norm": 0.017120331578900336, "grad_rms": 5.478513117749386e-05, "grad_to_weight": 0.0013301063841178273, "update_to_weight": 5.1092050814731415e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.039823442121637, "grad_norm": 0.05029526260848369, "grad_rms": 0.0011442568449222025, "grad_to_weight": 0.009979568369028155, "update_to_weight": 1.85389986286579e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.645654682956113, "grad_norm": 0.1094818870198591, "grad_rms": 0.0002638125418476059, "grad_to_weight": 0.0065772052289394325, "update_to_weight": 4.76342240161546e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57688579600124, "grad_norm": 0.03902004587338383, "grad_rms": 5.3834017831192726e-05, "grad_to_weight": 0.0015876725064870728, "update_to_weight": 6.403475254610088e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.834320327495362, "grad_norm": 0.025246849274445142, "grad_rms": 3.0649086859403084e-05, "grad_to_weight": 0.0008187905232317445, "update_to_weight": 5.7367605198711405e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.64074833649671, "grad_norm": 0.018888396373037727, "grad_rms": 1.6956155688195857e-05, "grad_to_weight": 0.0005452652520539614, "update_to_weight": 6.816851407178644e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09384708022258, "grad_norm": 0.014108938250733895, "grad_rms": 1.1348691180631466e-05, "grad_to_weight": 0.00038035791273470775, "update_to_weight": 7.120197513426271e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.67377733720432, "grad_norm": 0.004895709187186682, "grad_rms": 8.845035528336367e-06, "grad_to_weight": 0.00019841749888067174, "update_to_weight": 4.900609158750474e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.371266666593264, "grad_norm": 0.013399345016016626, "grad_rms": 1.0777921478364696e-05, "grad_to_weight": 0.00035854671813932655, "update_to_weight": 7.580798453441062e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.520027907050704, "grad_norm": 0.026102190292670337, "grad_rms": 1.158846809888598e-05, "grad_to_weight": 0.000410928507948168, "update_to_weight": 8.165321882218108e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.787917831471, "grad_norm": 0.033455297800421085, "grad_rms": 2.4963373270037327e-05, "grad_to_weight": 0.0006587255242759792, "update_to_weight": 6.116280044824617e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11542489433616, "grad_norm": 0.03316855032921657, "grad_rms": 2.8483841961203223e-05, "grad_to_weight": 0.0008268278453134268, "update_to_weight": 6.342990391819479e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205661243301407, "grad_norm": 0.060172144226783436, "grad_rms": 9.379984593084111e-05, "grad_to_weight": 0.0021333357054720277, "update_to_weight": 4.996636018939289e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.418623995669883, "grad_norm": 0.005634325940433079, "grad_rms": 0.00047962567532170867, "grad_to_weight": 0.0006692692230144847, "update_to_weight": 2.585548115644513e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3253281708744087, "grad_norm": 0.060589169911620576, "grad_rms": 0.0014026172730551565, "grad_to_weight": 0.045716352555643484, "update_to_weight": 6.843070675222958e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49850}
{"step": 49900, "lr": 1.0001004005696417e-05, "window_steps": 50, "elapsed_seconds": 5123.894989134999, "steps_per_second": 1.7369598750310205, "eta_seconds": 57.57185381050561, "images_seen": 51097600, "images_per_second": 1778.646912031765, "effective_batch_size": 1024, "loss": 0.3628781735897064, "edge_loss": 0.18825218081474304, "diag_loss": 0.17462590336799622, "jvp_rms": 2.2626171112060547, "grad_norm": 0.28342947363853455, "grad_norm_postclip": 0.2670843303203583, "clip_fraction": 0.07999999821186066, "ema_updates": 49900, "ema_decay": 0.9998000100000001, "loss_nfe_2": 0.34974223375320435, "updates_nfe_2": 8, "loss_nfe_4": 0.31124616662661236, "updates_nfe_4": 12, "loss_nfe_3": 0.3521909782519707, "updates_nfe_3": 13, "loss_nfe_5": 0.4190597001995359, "updates_nfe_5": 7, "loss_nfe_1": 0.4099114567041397, "updates_nfe_1": 10, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.51587722282571, "grad_norm": 0.012684712776244854, "grad_rms": 3.69366446153063e-05, "grad_to_weight": 0.000768031428491972, "update_to_weight": 4.654277837826262e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.26144616409181, "grad_norm": 0.004693029663761612, "grad_rms": 1.366564398557628e-05, "grad_to_weight": 0.0002885985426145347, "update_to_weight": 4.714179381843278e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067404549390891, "grad_norm": 0.0011016415212137703, "grad_rms": 1.9722546980015208e-05, "grad_to_weight": 0.00021739758696514354, "update_to_weight": 2.2948567976679037e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.872653191759362, "grad_norm": 0.006696542147470817, "grad_rms": 2.1428962300995913e-05, "grad_to_weight": 0.0005202146012725424, "update_to_weight": 5.246355514740088e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.039499852326288, "grad_norm": 0.010863906876114486, "grad_rms": 0.00024716243957924046, "grad_to_weight": 0.0021557510059454785, "update_to_weight": 1.6249587133383266e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.646004852187883, "grad_norm": 0.03557770487278029, "grad_rms": 8.572965822090377e-05, "grad_to_weight": 0.002137311936930266, "update_to_weight": 5.9318055423481577e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.57720367392693, "grad_norm": 0.019059407187757334, "grad_rms": 2.629531676428837e-05, "grad_to_weight": 0.0007754912821093952, "update_to_weight": 8.568615726748703e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.834055079355338, "grad_norm": 0.012680248805265514, "grad_rms": 1.5393526645909208e-05, "grad_to_weight": 0.000411241686266412, "update_to_weight": 8.254641220049744e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63984444598716, "grad_norm": 0.009550327163342325, "grad_rms": 8.573350064063426e-06, "grad_to_weight": 0.000275703523387175, "update_to_weight": 9.243627243456444e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09369076448514, "grad_norm": 0.007880287915614612, "grad_rms": 6.338602691391069e-06, "grad_to_weight": 0.00021244281044040752, "update_to_weight": 9.590036945662144e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673287542685966, "grad_norm": 0.002796767615259387, "grad_rms": 5.052895908567081e-06, "grad_to_weight": 0.00011335204562508525, "update_to_weight": 6.962671683184108e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.37145611856561, "grad_norm": 0.00806243283683688, "grad_rms": 6.485113110838432e-06, "grad_to_weight": 0.00021573772269557295, "update_to_weight": 9.100813898899865e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.519974097719846, "grad_norm": 0.018114629707115745, "grad_rms": 8.042267952624275e-06, "grad_to_weight": 0.0002851800550051861, "update_to_weight": 9.003003472633822e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78800916555009, "grad_norm": 0.023826379252801736, "grad_rms": 1.7778553415049154e-05, "grad_to_weight": 0.0004691339480375489, "update_to_weight": 6.252351348254964e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.115698761325035, "grad_norm": 0.023663215007211937, "grad_rms": 2.0321035133262555e-05, "grad_to_weight": 0.0005898741823743402, "update_to_weight": 6.52176266509077e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.205876553795818, "grad_norm": 0.04721623137331676, "grad_rms": 7.360341375836671e-05, "grad_to_weight": 0.0016739856066257444, "update_to_weight": 4.489864727106729e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.4184232584666, "grad_norm": 0.005183597066680757, "grad_rms": 0.0004412570855833887, "grad_to_weight": 0.0006157444105067413, "update_to_weight": 3.303534463456521e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3250922188966028, "grad_norm": 0.07208788122638901, "grad_rms": 0.0016688082628227138, "grad_to_weight": 0.05440216175023366, "update_to_weight": 5.1849152977826646e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49900}
{"step": 49950, "lr": 1.0000255997064122e-05, "window_steps": 50, "elapsed_seconds": 5151.369800064, "steps_per_second": 1.7374019624622574, "eta_seconds": 28.778602234994413, "images_seen": 51148800, "images_per_second": 1779.0996095613516, "effective_batch_size": 1024, "loss": 0.36172327399253845, "edge_loss": 0.18886211514472961, "diag_loss": 0.17286120355129242, "jvp_rms": 2.259570360183716, "grad_norm": 0.19331985712051392, "grad_norm_postclip": 0.19331985712051392, "clip_fraction": 0.0, "ema_updates": 49950, "ema_decay": 0.9998000100000001, "loss_nfe_1": 0.40308638768536703, "updates_nfe_1": 14, "loss_nfe_5": 0.3075896731831811, "updates_nfe_5": 11, "loss_nfe_3": 0.35274266203244525, "updates_nfe_3": 6, "loss_nfe_4": 0.37225793302059174, "updates_nfe_4": 12, "loss_nfe_2": 0.35370273675237385, "updates_nfe_2": 7, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.513395478356635, "grad_norm": 0.01348912068164534, "grad_rms": 3.9279001864668984e-05, "grad_to_weight": 0.0008168593006401938, "update_to_weight": 3.338760618962932e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262376972175876, "grad_norm": 0.011030507334833612, "grad_rms": 3.2119759945711044e-05, "grad_to_weight": 0.000678283829830428, "update_to_weight": 3.442994510912784e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067506739758233, "grad_norm": 0.000639440095929676, "grad_rms": 1.1447814093811071e-05, "grad_to_weight": 0.0001261843602323819, "update_to_weight": 1.8090803135181045e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.871756432105952, "grad_norm": 0.006788913472880098, "grad_rms": 2.1724550920659286e-05, "grad_to_weight": 0.0005274271237720556, "update_to_weight": 4.118825547526462e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.039405231501062, "grad_norm": 0.01817187889669808, "grad_rms": 0.00041342455997310436, "grad_to_weight": 0.003605957064755698, "update_to_weight": 1.8315268893314373e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64559143170732, "grad_norm": 0.06209601095061674, "grad_rms": 0.00014962937645117012, "grad_to_weight": 0.0037304778989308415, "update_to_weight": 5.025701026708174e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.576822468267103, "grad_norm": 0.016333649035412828, "grad_rms": 2.253472372313755e-05, "grad_to_weight": 0.0006645956391027511, "update_to_weight": 6.126389317271075e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.833869008402818, "grad_norm": 0.008093130857177805, "grad_rms": 9.82487232009703e-06, "grad_to_weight": 0.00026247535964339316, "update_to_weight": 4.835783895388619e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.640040026635, "grad_norm": 0.004774075148787114, "grad_rms": 4.2856979433962566e-06, "grad_to_weight": 0.0001378195621343477, "update_to_weight": 5.560106054279697e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09467801071671, "grad_norm": 0.004113210261939494, "grad_rms": 3.3085092722216346e-06, "grad_to_weight": 0.0001108841074385706, "update_to_weight": 5.9690719384139016e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673145302557064, "grad_norm": 0.0013722353178914437, "grad_rms": 2.479205703589302e-06, "grad_to_weight": 5.561655399278294e-05, "update_to_weight": 3.7338092697602883e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.37213389146467, "grad_norm": 0.00428337784340187, "grad_rms": 3.4453855769195443e-06, "grad_to_weight": 0.00011461421645982436, "update_to_weight": 6.35515082055131e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.52163447893167, "grad_norm": 0.011415345646073723, "grad_rms": 5.068017946924139e-06, "grad_to_weight": 0.00017970799617663918, "update_to_weight": 7.445425291799872e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.78822067541843, "grad_norm": 0.02284985957769672, "grad_rms": 1.704990274511325e-05, "grad_to_weight": 0.0004499047077023528, "update_to_weight": 6.204449353617481e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.11614759484717, "grad_norm": 0.03229123187028769, "grad_rms": 2.7730435493759138e-05, "grad_to_weight": 0.0008049434904969646, "update_to_weight": 7.002214845998794e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.20563094822557, "grad_norm": 0.0630336937839606, "grad_rms": 9.826059618389942e-05, "grad_to_weight": 0.002234791127334313, "update_to_weight": 5.048773665908777e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.418751776227078, "grad_norm": 0.006688398039581478, "grad_rms": 0.0005693542511507744, "grad_to_weight": 0.0007944643359682152, "update_to_weight": 4.254284627585354e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.32555535751779, "grad_norm": 0.07982105185375486, "grad_rms": 0.0018478283535955007, "grad_to_weight": 0.06021706404116253, "update_to_weight": 9.916448460741668e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 49950}
{"step": 50000, "lr": 1.0000000098422651e-05, "window_steps": 50, "elapsed_seconds": 5178.840372322999, "steps_per_second": 1.7378407815190096, "eta_seconds": 0.0, "images_seen": 51200000, "images_per_second": 1779.5489602754658, "effective_batch_size": 1024, "loss": 0.3548988401889801, "edge_loss": 0.18379321694374084, "diag_loss": 0.17110557854175568, "jvp_rms": 1.8925390243530273, "grad_norm": 0.21770034730434418, "grad_norm_postclip": 0.21770034730434418, "clip_fraction": 0.0, "ema_updates": 50000, "ema_decay": 0.9998000100000001, "loss_nfe_3": 0.35597619265317915, "updates_nfe_3": 10, "loss_nfe_4": 0.37776655620998806, "updates_nfe_4": 9, "loss_nfe_5": 0.3045877367258072, "updates_nfe_5": 8, "loss_nfe_2": 0.3413205325603485, "updates_nfe_2": 15, "loss_nfe_1": 0.403596218675375, "updates_nfe_1": 8, "gpu_allocated_gb": 0.8975429534912109, "gpu_reserved_gb": 43.775390625, "gpu_peak_allocated_gb": 37.954071044921875, "gradient_groups": {"t_emb": {"weight_norm": 16.515023530024582, "grad_norm": 0.11432138125518912, "grad_rms": 0.00033289269578588913, "grad_to_weight": 0.00692226572050297, "update_to_weight": 5.734569434715802e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "r_emb": {"weight_norm": 16.262888330958752, "grad_norm": 0.039618767773493664, "grad_rms": 0.00011536598196266779, "grad_to_weight": 0.002436145841207901, "update_to_weight": 4.733992422527738e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "y_emb": {"weight_norm": 5.067645060453645, "grad_norm": 0.0034476322407670404, "grad_rms": 6.172251819577942e-05, "grad_to_weight": 0.000680322358736469, "update_to_weight": 2.5612938476024494e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "cond_out": {"weight_norm": 12.87244278780043, "grad_norm": 0.03697475371479452, "grad_rms": 0.00011831936333622446, "grad_to_weight": 0.002872396042019042, "update_to_weight": 6.601609987986573e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "in_conv": {"weight_norm": 5.039121547619157, "grad_norm": 0.02010702286093248, "grad_rms": 0.000457450609587797, "grad_to_weight": 0.0039901841364458615, "update_to_weight": 3.290046528251253e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.0": {"weight_norm": 16.64608276842342, "grad_norm": 0.09815406752986451, "grad_rms": 0.0002365165120239306, "grad_to_weight": 0.0058965264618326095, "update_to_weight": 8.954140929857495e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.1": {"weight_norm": 24.577499197373093, "grad_norm": 0.01800084380436817, "grad_rms": 2.4834869479276555e-05, "grad_to_weight": 0.0007324115305552383, "update_to_weight": 9.463347569559213e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.2": {"weight_norm": 30.83386603917639, "grad_norm": 0.008038267314234277, "grad_rms": 9.758269257084616e-06, "grad_to_weight": 0.0002606960575109572, "update_to_weight": 6.642526625302697e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "down_levels.3": {"weight_norm": 34.63907268454675, "grad_norm": 0.003863454481162671, "grad_rms": 3.468231732491786e-06, "grad_to_weight": 0.00011153458166581472, "update_to_weight": 7.572530100937947e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid1": {"weight_norm": 37.09438513040922, "grad_norm": 0.0032076774317542586, "grad_rms": 2.580133241292384e-06, "grad_to_weight": 8.647339537984874e-05, "update_to_weight": 7.81818486024589e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid_attn": {"weight_norm": 24.673390202443425, "grad_norm": 0.0011791270861310043, "grad_rms": 2.1303187281934394e-06, "grad_to_weight": 4.778942319869097e-05, "update_to_weight": 5.262799180821532e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "mid2": {"weight_norm": 37.37197710653749, "grad_norm": 0.003229066174170447, "grad_rms": 2.597337528965184e-06, "grad_to_weight": 8.640340769141662e-05, "update_to_weight": 7.984646921565486e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.0": {"weight_norm": 63.52169789540913, "grad_norm": 0.008797741452653435, "grad_rms": 3.905892380033334e-06, "grad_to_weight": 0.0001384997842334008, "update_to_weight": 9.197308783625432e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.1": {"weight_norm": 50.788316815976486, "grad_norm": 0.022728171192384107, "grad_rms": 1.695910240002863e-05, "grad_to_weight": 0.00044750786435266354, "update_to_weight": 8.085612339796532e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.2": {"weight_norm": 40.116334734120464, "grad_norm": 0.038615167286647545, "grad_rms": 3.316118164288409e-05, "grad_to_weight": 0.0009625796459865482, "update_to_weight": 9.42659917428506e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "up_levels.3": {"weight_norm": 28.206099436294913, "grad_norm": 0.10106223952213236, "grad_rms": 0.00015754171001242578, "grad_to_weight": 0.0035829923861109263, "update_to_weight": 8.058384138680281e-05, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "out_norm": {"weight_norm": 8.418451806035767, "grad_norm": 0.00825520231395145, "grad_rms": 0.0007027294882485905, "grad_to_weight": 0.0009806081336752117, "update_to_weight": 5.2765190801682055e-06, "missing_grad_tensors": 0, "zero_grad_tensors": 0}, "u_head": {"weight_norm": 1.3251545617986, "grad_norm": 0.10927559814558875, "grad_rms": 0.002529690400215243, "grad_to_weight": 0.08246253025554366, "update_to_weight": 0.00010683296531524654, "missing_grad_tensors": 0, "zero_grad_tensors": 0}}, "gradient_snapshot_step": 50000}