diff --git a/examples/huggingface/llama_B_48_no_patching.log b/examples/huggingface/llama_B_48_no_patching.log
index 5c002fb4a..d33cd63ed 100644
--- a/examples/huggingface/llama_B_48_no_patching.log
+++ b/examples/huggingface/llama_B_48_no_patching.log
@@ -1,21 +1,21 @@
-{'loss': 1.3646, 'grad_norm': 12.249734878540039, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 52992}
-{'loss': 1.2759, 'grad_norm': 14.460180282592773, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 117504, 'step': 2, 'step_time_sec': 6.35, 'avg_step_time_sec': 6.35, 'time_to_completion_sec': 114.29, 'estimated_total_time_sec': 126.98, 'step_peak_memory_allocated_MB': 52216.61, 'step_peak_memory_reserved_MB': 75474.0, 'total_peak_memory_allocated_MB': 52216.61, 'total_peak_memory_reserved_MB': 75474.0, 'step_tokens_per_second': 10160.61, 'avg_tokens_per_second': 10160.61}
-{'loss': 1.3749, 'grad_norm': 8.127913475036621, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 174336, 'step': 3, 'step_time_sec': 7.83, 'avg_step_time_sec': 7.09, 'time_to_completion_sec': 120.56, 'estimated_total_time_sec': 141.84, 'step_peak_memory_allocated_MB': 64316.86, 'step_peak_memory_reserved_MB': 79572.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 7254.2, 'avg_tokens_per_second': 8555.24}
-{'loss': 1.3761, 'grad_norm': 12.362775802612305, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 228096, 'step': 4, 'step_time_sec': 7.68, 'avg_step_time_sec': 7.29, 'time_to_completion_sec': 116.59, 'estimated_total_time_sec': 145.73, 'step_peak_memory_allocated_MB': 46839.06, 'step_peak_memory_reserved_MB': 68064.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 7003.12, 'avg_tokens_per_second': 8010.19}
-{'loss': 1.3788, 'grad_norm': 5.331186771392822, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 284160, 'step': 5, 'step_time_sec': 5.33, 'avg_step_time_sec': 6.8, 'time_to_completion_sec': 101.95, 'estimated_total_time_sec': 135.94, 'step_peak_memory_allocated_MB': 52217.09, 'step_peak_memory_reserved_MB': 75556.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 10523.68, 'avg_tokens_per_second': 8502.71}
-{'loss': 1.2772, 'grad_norm': 5.286765098571777, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 331008, 'step': 6, 'step_time_sec': 4.61, 'avg_step_time_sec': 6.36, 'time_to_completion_sec': 89.02, 'estimated_total_time_sec': 127.18, 'step_peak_memory_allocated_MB': 42807.18, 'step_peak_memory_reserved_MB': 75556.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 10169.6, 'avg_tokens_per_second': 8744.22}
-{'loss': 1.311, 'grad_norm': 4.666638374328613, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 386304, 'step': 7, 'step_time_sec': 6.12, 'avg_step_time_sec': 6.32, 'time_to_completion_sec': 82.15, 'estimated_total_time_sec': 126.38, 'step_peak_memory_allocated_MB': 56248.4, 'step_peak_memory_reserved_MB': 75556.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 9036.85, 'avg_tokens_per_second': 8791.45}
-{'loss': 1.2994, 'grad_norm': 3.4744791984558105, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 446208, 'step': 8, 'step_time_sec': 6.84, 'avg_step_time_sec': 6.39, 'time_to_completion_sec': 76.72, 'estimated_total_time_sec': 127.87, 'step_peak_memory_allocated_MB': 58937.21, 'step_peak_memory_reserved_MB': 78436.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 8756.77, 'avg_tokens_per_second': 8786.15}
-{'loss': 1.1939, 'grad_norm': 3.459392786026001, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 503040, 'step': 9, 'step_time_sec': 5.23, 'avg_step_time_sec': 6.25, 'time_to_completion_sec': 68.73, 'estimated_total_time_sec': 124.96, 'step_peak_memory_allocated_MB': 46839.06, 'step_peak_memory_reserved_MB': 67228.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 10867.5, 'avg_tokens_per_second': 9003.91}
-{'loss': 1.2519, 'grad_norm': 3.662140130996704, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 562176, 'step': 10, 'step_time_sec': 6.48, 'avg_step_time_sec': 6.27, 'time_to_completion_sec': 62.74, 'estimated_total_time_sec': 125.47, 'step_peak_memory_allocated_MB': 50871.82, 'step_peak_memory_reserved_MB': 74344.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 9127.49, 'avg_tokens_per_second': 9018.09}
-{'loss': 1.1777, 'grad_norm': 3.343947649002075, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 622080, 'step': 11, 'step_time_sec': 5.96, 'avg_step_time_sec': 6.24, 'time_to_completion_sec': 56.18, 'estimated_total_time_sec': 124.85, 'step_peak_memory_allocated_MB': 56248.4, 'step_peak_memory_reserved_MB': 77294.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 10050.35, 'avg_tokens_per_second': 9116.66}
-{'loss': 1.231, 'grad_norm': 3.3978335857391357, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 665088, 'step': 12, 'step_time_sec': 4.39, 'avg_step_time_sec': 6.07, 'time_to_completion_sec': 48.59, 'estimated_total_time_sec': 121.48, 'step_peak_memory_allocated_MB': 42807.18, 'step_peak_memory_reserved_MB': 71576.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 9798.88, 'avg_tokens_per_second': 9161.47}
-{'loss': 1.1562, 'grad_norm': 3.410357713699341, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 721152, 'step': 13, 'step_time_sec': 5.25, 'avg_step_time_sec': 6.01, 'time_to_completion_sec': 42.04, 'estimated_total_time_sec': 120.11, 'step_peak_memory_allocated_MB': 53559.79, 'step_peak_memory_reserved_MB': 71576.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 10671.97, 'avg_tokens_per_second': 9271.58}
-{'loss': 1.2285, 'grad_norm': 3.1715331077575684, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 784128, 'step': 14, 'step_time_sec': 8.31, 'avg_step_time_sec': 6.18, 'time_to_completion_sec': 37.1, 'estimated_total_time_sec': 123.65, 'step_peak_memory_allocated_MB': 52216.12, 'step_peak_memory_reserved_MB': 71576.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 7578.9, 'avg_tokens_per_second': 9096.59}
-{'loss': 1.2367, 'grad_norm': 3.1703574657440186, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 842496, 'step': 15, 'step_time_sec': 6.45, 'avg_step_time_sec': 6.2, 'time_to_completion_sec': 31.01, 'estimated_total_time_sec': 124.03, 'step_peak_memory_allocated_MB': 45495.77, 'step_peak_memory_reserved_MB': 71576.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 9055.83, 'avg_tokens_per_second': 9093.56}
-{'loss': 1.122, 'grad_norm': 3.0725834369659424, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 890880, 'step': 16, 'step_time_sec': 4.02, 'avg_step_time_sec': 6.06, 'time_to_completion_sec': 24.22, 'estimated_total_time_sec': 121.11, 'step_peak_memory_allocated_MB': 48183.11, 'step_peak_memory_reserved_MB': 71576.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 12049.45, 'avg_tokens_per_second': 9224.23}
-{'loss': 1.2217, 'grad_norm': 3.2708375453948975, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 943872, 'step': 17, 'step_time_sec': 5.22, 'avg_step_time_sec': 6.0, 'time_to_completion_sec': 18.01, 'estimated_total_time_sec': 120.06, 'step_peak_memory_allocated_MB': 45495.77, 'step_peak_memory_reserved_MB': 71576.0, 'total_peak_memory_allocated_MB': 64316.86, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 10161.25, 'avg_tokens_per_second': 9275.11}
-{'loss': 1.2064, 'grad_norm': 3.1343138217926025, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1012224, 'step': 18, 'step_time_sec': 8.62, 'avg_step_time_sec': 6.16, 'time_to_completion_sec': 12.31, 'estimated_total_time_sec': 123.14, 'step_peak_memory_allocated_MB': 67004.76, 'step_peak_memory_reserved_MB': 79530.0, 'total_peak_memory_allocated_MB': 67004.76, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 7927.14, 'avg_tokens_per_second': 9164.07}
-{'loss': 1.2952, 'grad_norm': 3.0990655422210693, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1068288, 'step': 19, 'step_time_sec': 4.43, 'avg_step_time_sec': 6.06, 'time_to_completion_sec': 6.06, 'estimated_total_time_sec': 121.23, 'step_peak_memory_allocated_MB': 50871.59, 'step_peak_memory_reserved_MB': 68464.0, 'total_peak_memory_allocated_MB': 67004.76, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 12655.61, 'avg_tokens_per_second': 9305.84}
-{'loss': 1.1813, 'grad_norm': 3.002934217453003, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1122048, 'step': 20, 'step_time_sec': 5.84, 'avg_step_time_sec': 6.05, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 120.99, 'step_peak_memory_allocated_MB': 48183.11, 'step_peak_memory_reserved_MB': 68464.0, 'total_peak_memory_allocated_MB': 67004.76, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 9205.12, 'avg_tokens_per_second': 9300.72}
-{'train_runtime': 123.166, 'train_samples_per_second': 31.177, 'train_steps_per_second': 0.162, 'train_loss': 1.2580344438552857, 'epoch': 0.08, 'num_input_tokens_seen': 1122048, 'step': 20, 'step_time_sec': 5.84, 'avg_step_time_sec': 6.05, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 120.99, 'step_peak_memory_allocated_MB': 48183.11, 'step_peak_memory_reserved_MB': 68464.0, 'total_peak_memory_allocated_MB': 67004.76, 'total_peak_memory_reserved_MB': 79572.0, 'step_tokens_per_second': 9205.12, 'avg_tokens_per_second': 9300.72}
+{'loss': 1.3883, 'grad_norm': 11.805442810058594, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 54528}
+{'loss': 1.403, 'grad_norm': 11.622239112854004, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 125184, 'step': 2, 'step_time_sec': 8.97, 'avg_step_time_sec': 8.97, 'time_to_completion_sec': 161.37, 'estimated_total_time_sec': 179.3, 'step_peak_memory_allocated_MB': 56248.4, 'step_peak_memory_reserved_MB': 74236.0, 'total_peak_memory_allocated_MB': 56248.4, 'total_peak_memory_reserved_MB': 74236.0, 'step_tokens_per_second': 7881.2, 'avg_tokens_per_second': 7881.2}
+{'loss': 1.3964, 'grad_norm': 8.47226333618164, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 185856, 'step': 3, 'step_time_sec': 6.5, 'avg_step_time_sec': 7.73, 'time_to_completion_sec': 131.47, 'estimated_total_time_sec': 154.67, 'step_peak_memory_allocated_MB': 52216.12, 'step_peak_memory_reserved_MB': 74738.0, 'total_peak_memory_allocated_MB': 56248.4, 'total_peak_memory_reserved_MB': 74738.0, 'step_tokens_per_second': 9331.88, 'avg_tokens_per_second': 8491.0}
+{'loss': 1.2832, 'grad_norm': 9.746687889099121, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 254208, 'step': 4, 'step_time_sec': 9.62, 'avg_step_time_sec': 8.36, 'time_to_completion_sec': 133.8, 'estimated_total_time_sec': 167.25, 'step_peak_memory_allocated_MB': 68349.51, 'step_peak_memory_reserved_MB': 78662.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 7104.96, 'avg_tokens_per_second': 7959.49}
+{'loss': 1.2847, 'grad_norm': 4.774288654327393, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 308736, 'step': 5, 'step_time_sec': 5.55, 'avg_step_time_sec': 7.66, 'time_to_completion_sec': 114.89, 'estimated_total_time_sec': 153.19, 'step_peak_memory_allocated_MB': 56248.4, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 9823.21, 'avg_tokens_per_second': 8297.15}
+{'loss': 1.3723, 'grad_norm': 5.747477054595947, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 359424, 'step': 6, 'step_time_sec': 6.64, 'avg_step_time_sec': 7.46, 'time_to_completion_sec': 104.39, 'estimated_total_time_sec': 149.12, 'step_peak_memory_allocated_MB': 44151.09, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 7630.72, 'avg_tokens_per_second': 8178.41}
+{'loss': 1.2954, 'grad_norm': 4.253581523895264, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 413184, 'step': 7, 'step_time_sec': 4.75, 'avg_step_time_sec': 7.01, 'time_to_completion_sec': 91.07, 'estimated_total_time_sec': 140.11, 'step_peak_memory_allocated_MB': 46839.06, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 11314.82, 'avg_tokens_per_second': 8532.95}
+{'loss': 1.2282, 'grad_norm': 3.626068115234375, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 461568, 'step': 8, 'step_time_sec': 5.74, 'avg_step_time_sec': 6.82, 'time_to_completion_sec': 81.9, 'estimated_total_time_sec': 136.5, 'step_peak_memory_allocated_MB': 41463.32, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 8425.85, 'avg_tokens_per_second': 8520.08}
+{'loss': 1.2229, 'grad_norm': 4.115623950958252, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 509952, 'step': 9, 'step_time_sec': 3.66, 'avg_step_time_sec': 6.43, 'time_to_completion_sec': 70.73, 'estimated_total_time_sec': 128.59, 'step_peak_memory_allocated_MB': 45495.05, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 13206.87, 'avg_tokens_per_second': 8853.88}
+{'loss': 1.2422, 'grad_norm': 3.959798574447632, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 568320, 'step': 10, 'step_time_sec': 6.19, 'avg_step_time_sec': 6.4, 'time_to_completion_sec': 64.04, 'estimated_total_time_sec': 128.07, 'step_peak_memory_allocated_MB': 45495.05, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 9421.98, 'avg_tokens_per_second': 8914.95}
+{'loss': 1.2622, 'grad_norm': 3.6134583950042725, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 616704, 'step': 11, 'step_time_sec': 4.82, 'avg_step_time_sec': 6.25, 'time_to_completion_sec': 56.21, 'estimated_total_time_sec': 124.9, 'step_peak_memory_allocated_MB': 46839.06, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 10038.49, 'avg_tokens_per_second': 9001.66}
+{'loss': 1.2243, 'grad_norm': 3.3471832275390625, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 669696, 'step': 12, 'step_time_sec': 5.3, 'avg_step_time_sec': 6.16, 'time_to_completion_sec': 49.28, 'estimated_total_time_sec': 123.19, 'step_peak_memory_allocated_MB': 40119.5, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 9992.91, 'avg_tokens_per_second': 9079.24}
+{'loss': 1.1169, 'grad_norm': 3.6069729328155518, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 717312, 'step': 13, 'step_time_sec': 6.13, 'avg_step_time_sec': 6.16, 'time_to_completion_sec': 43.1, 'estimated_total_time_sec': 123.14, 'step_peak_memory_allocated_MB': 41463.32, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 7772.04, 'avg_tokens_per_second': 8970.84}
+{'loss': 1.2609, 'grad_norm': 2.9739620685577393, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 777216, 'step': 14, 'step_time_sec': 5.09, 'avg_step_time_sec': 6.07, 'time_to_completion_sec': 36.45, 'estimated_total_time_sec': 121.49, 'step_peak_memory_allocated_MB': 54904.68, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 11774.2, 'avg_tokens_per_second': 9151.45}
+{'loss': 1.2699, 'grad_norm': 3.1643807888031006, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 840960, 'step': 15, 'step_time_sec': 8.42, 'avg_step_time_sec': 6.24, 'time_to_completion_sec': 31.21, 'estimated_total_time_sec': 124.85, 'step_peak_memory_allocated_MB': 42807.18, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 7567.31, 'avg_tokens_per_second': 8998.76}
+{'loss': 1.1884, 'grad_norm': 3.0365970134735107, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 901632, 'step': 16, 'step_time_sec': 6.66, 'avg_step_time_sec': 6.27, 'time_to_completion_sec': 25.08, 'estimated_total_time_sec': 125.41, 'step_peak_memory_allocated_MB': 49527.21, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 9103.53, 'avg_tokens_per_second': 9006.18}
+{'loss': 1.2402, 'grad_norm': 3.1069788932800293, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 946944, 'step': 17, 'step_time_sec': 4.2, 'avg_step_time_sec': 6.14, 'time_to_completion_sec': 18.42, 'estimated_total_time_sec': 122.82, 'step_peak_memory_allocated_MB': 49527.21, 'step_peak_memory_reserved_MB': 75870.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 10792.52, 'avg_tokens_per_second': 9082.51}
+{'loss': 1.1933, 'grad_norm': 2.8866586685180664, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1015296, 'step': 18, 'step_time_sec': 8.11, 'avg_step_time_sec': 6.26, 'time_to_completion_sec': 12.51, 'estimated_total_time_sec': 125.13, 'step_peak_memory_allocated_MB': 60281.83, 'step_peak_memory_reserved_MB': 77852.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 8430.29, 'avg_tokens_per_second': 9032.8}
+{'loss': 1.192, 'grad_norm': 3.2670505046844482, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1065984, 'step': 19, 'step_time_sec': 4.61, 'avg_step_time_sec': 6.17, 'time_to_completion_sec': 6.17, 'estimated_total_time_sec': 123.31, 'step_peak_memory_allocated_MB': 41464.29, 'step_peak_memory_reserved_MB': 65804.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 10988.89, 'avg_tokens_per_second': 9114.1}
+{'loss': 1.2296, 'grad_norm': 2.858172655105591, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1122816, 'step': 20, 'step_time_sec': 5.63, 'avg_step_time_sec': 6.14, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 122.74, 'step_peak_memory_allocated_MB': 50871.82, 'step_peak_memory_reserved_MB': 72944.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 10095.04, 'avg_tokens_per_second': 9161.46}
+{'train_runtime': 126.104, 'train_samples_per_second': 30.451, 'train_steps_per_second': 0.159, 'train_loss': 1.2647102057933808, 'epoch': 0.08, 'num_input_tokens_seen': 1122816, 'step': 20, 'step_time_sec': 5.63, 'avg_step_time_sec': 6.14, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 122.74, 'step_peak_memory_allocated_MB': 50871.82, 'step_peak_memory_reserved_MB': 72944.0, 'total_peak_memory_allocated_MB': 68349.51, 'total_peak_memory_reserved_MB': 78662.0, 'step_tokens_per_second': 10095.04, 'avg_tokens_per_second': 9161.46}
diff --git a/examples/huggingface/llama_B_48_post_init_class_patching.log b/examples/huggingface/llama_B_48_post_init_class_patching.log
index 3651974db..ac9db949a 100644
--- a/examples/huggingface/llama_B_48_post_init_class_patching.log
+++ b/examples/huggingface/llama_B_48_post_init_class_patching.log
@@ -1,21 +1,21 @@
-{'loss': 1.4339, 'grad_norm': 15.349852561950684, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 44544}
-{'loss': 1.3817, 'grad_norm': 13.617855072021484, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 102912, 'step': 2, 'step_time_sec': 4.57, 'avg_step_time_sec': 4.57, 'time_to_completion_sec': 82.24, 'estimated_total_time_sec': 91.38, 'step_peak_memory_allocated_MB': 38360.39, 'step_peak_memory_reserved_MB': 56552.0, 'total_peak_memory_allocated_MB': 38360.39, 'total_peak_memory_reserved_MB': 56552.0, 'step_tokens_per_second': 12775.05, 'avg_tokens_per_second': 12775.05}
-{'loss': 1.287, 'grad_norm': 8.358651161193848, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 155904, 'step': 3, 'step_time_sec': 4.57, 'avg_step_time_sec': 4.57, 'time_to_completion_sec': 77.68, 'estimated_total_time_sec': 91.39, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 56552.0, 'total_peak_memory_allocated_MB': 38360.39, 'total_peak_memory_reserved_MB': 56552.0, 'step_tokens_per_second': 11594.48, 'avg_tokens_per_second': 12184.67}
-{'loss': 1.3154, 'grad_norm': 12.276527404785156, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 218112, 'step': 4, 'step_time_sec': 6.44, 'avg_step_time_sec': 5.19, 'time_to_completion_sec': 83.08, 'estimated_total_time_sec': 103.85, 'step_peak_memory_allocated_MB': 38307.78, 'step_peak_memory_reserved_MB': 58558.0, 'total_peak_memory_allocated_MB': 38360.39, 'total_peak_memory_reserved_MB': 58558.0, 'step_tokens_per_second': 9663.36, 'avg_tokens_per_second': 11142.68}
-{'loss': 1.2961, 'grad_norm': 4.9072442054748535, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 285696, 'step': 5, 'step_time_sec': 7.47, 'avg_step_time_sec': 5.76, 'time_to_completion_sec': 86.43, 'estimated_total_time_sec': 115.23, 'step_peak_memory_allocated_MB': 40533.83, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 9047.51, 'avg_tokens_per_second': 10463.59}
-{'loss': 1.4121, 'grad_norm': 5.200970649719238, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 354048, 'step': 6, 'step_time_sec': 5.9, 'avg_step_time_sec': 5.79, 'time_to_completion_sec': 81.05, 'estimated_total_time_sec': 115.79, 'step_peak_memory_allocated_MB': 38307.89, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11584.08, 'avg_tokens_per_second': 10691.99}
-{'loss': 1.328, 'grad_norm': 4.505997180938721, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 403200, 'step': 7, 'step_time_sec': 4.59, 'avg_step_time_sec': 5.59, 'time_to_completion_sec': 72.66, 'estimated_total_time_sec': 111.79, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 10708.91, 'avg_tokens_per_second': 10694.31}
-{'loss': 1.3258, 'grad_norm': 3.3744163513183594, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 465408, 'step': 8, 'step_time_sec': 5.54, 'avg_step_time_sec': 5.58, 'time_to_completion_sec': 66.99, 'estimated_total_time_sec': 111.64, 'step_peak_memory_allocated_MB': 38360.39, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11232.44, 'avg_tokens_per_second': 10770.58}
-{'loss': 1.2751, 'grad_norm': 4.1117072105407715, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 523776, 'step': 9, 'step_time_sec': 5.7, 'avg_step_time_sec': 5.6, 'time_to_completion_sec': 61.56, 'estimated_total_time_sec': 111.93, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 10243.86, 'avg_tokens_per_second': 10703.55}
-{'loss': 1.2382, 'grad_norm': 3.717971086502075, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 585984, 'step': 10, 'step_time_sec': 6.5, 'avg_step_time_sec': 5.7, 'time_to_completion_sec': 56.97, 'estimated_total_time_sec': 113.94, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 9570.48, 'avg_tokens_per_second': 10559.91}
-{'loss': 1.1785, 'grad_norm': 3.4707894325256348, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 645888, 'step': 11, 'step_time_sec': 6.92, 'avg_step_time_sec': 5.82, 'time_to_completion_sec': 52.37, 'estimated_total_time_sec': 116.39, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 8655.49, 'avg_tokens_per_second': 10333.42}
-{'loss': 1.2379, 'grad_norm': 3.103318452835083, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 710400, 'step': 12, 'step_time_sec': 5.69, 'avg_step_time_sec': 5.81, 'time_to_completion_sec': 46.46, 'estimated_total_time_sec': 116.14, 'step_peak_memory_allocated_MB': 39223.21, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11346.5, 'avg_tokens_per_second': 10423.59}
-{'loss': 1.2656, 'grad_norm': 3.1070923805236816, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 778752, 'step': 13, 'step_time_sec': 5.85, 'avg_step_time_sec': 5.81, 'time_to_completion_sec': 40.68, 'estimated_total_time_sec': 116.22, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11676.28, 'avg_tokens_per_second': 10528.75}
-{'loss': 1.1827, 'grad_norm': 3.301771402359009, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 831744, 'step': 14, 'step_time_sec': 5.54, 'avg_step_time_sec': 5.79, 'time_to_completion_sec': 34.74, 'estimated_total_time_sec': 115.8, 'step_peak_memory_allocated_MB': 39223.21, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 9566.51, 'avg_tokens_per_second': 10457.94}
-{'loss': 1.1968, 'grad_norm': 3.0512115955352783, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 883200, 'step': 15, 'step_time_sec': 4.99, 'avg_step_time_sec': 5.73, 'time_to_completion_sec': 28.67, 'estimated_total_time_sec': 114.67, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 10305.86, 'avg_tokens_per_second': 10448.48}
-{'loss': 1.2275, 'grad_norm': 3.0405077934265137, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 940032, 'step': 16, 'step_time_sec': 5.19, 'avg_step_time_sec': 5.7, 'time_to_completion_sec': 22.79, 'estimated_total_time_sec': 113.94, 'step_peak_memory_allocated_MB': 39006.08, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 10952.44, 'avg_tokens_per_second': 10479.08}
-{'loss': 1.2437, 'grad_norm': 3.1204562187194824, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 986112, 'step': 17, 'step_time_sec': 3.95, 'avg_step_time_sec': 5.59, 'time_to_completion_sec': 16.76, 'estimated_total_time_sec': 111.76, 'step_peak_memory_allocated_MB': 38307.78, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11651.23, 'avg_tokens_per_second': 10530.93}
-{'loss': 1.2343, 'grad_norm': 2.86057710647583, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1040640, 'step': 18, 'step_time_sec': 4.66, 'avg_step_time_sec': 5.53, 'time_to_completion_sec': 11.07, 'estimated_total_time_sec': 110.67, 'step_peak_memory_allocated_MB': 38360.39, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11705.94, 'avg_tokens_per_second': 10589.11}
-{'loss': 1.2195, 'grad_norm': 2.9343435764312744, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1088256, 'step': 19, 'step_time_sec': 3.79, 'avg_step_time_sec': 5.44, 'time_to_completion_sec': 5.44, 'estimated_total_time_sec': 108.73, 'step_peak_memory_allocated_MB': 38307.78, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 12576.85, 'avg_tokens_per_second': 10666.02}
-{'loss': 1.2748, 'grad_norm': 3.041419506072998, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1141248, 'step': 20, 'step_time_sec': 4.6, 'avg_step_time_sec': 5.39, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 107.84, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11527.69, 'avg_tokens_per_second': 10704.68}
-{'train_runtime': 109.5837, 'train_samples_per_second': 35.042, 'train_steps_per_second': 0.183, 'train_loss': 1.27772376537323, 'epoch': 0.08, 'num_input_tokens_seen': 1141248, 'step': 20, 'step_time_sec': 4.6, 'avg_step_time_sec': 5.39, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 107.84, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 63122.0, 'total_peak_memory_allocated_MB': 40533.83, 'total_peak_memory_reserved_MB': 63122.0, 'step_tokens_per_second': 11527.69, 'avg_tokens_per_second': 10704.68}
+{'loss': 1.4012, 'grad_norm': 13.653097152709961, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 47616}
+{'loss': 1.3926, 'grad_norm': 11.953817367553711, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 125952, 'step': 2, 'step_time_sec': 6.81, 'avg_step_time_sec': 6.81, 'time_to_completion_sec': 122.55, 'estimated_total_time_sec': 136.17, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 56776.0, 'total_peak_memory_allocated_MB': 38307.84, 'total_peak_memory_reserved_MB': 56776.0, 'step_tokens_per_second': 11505.58, 'avg_tokens_per_second': 11505.58}
+{'loss': 1.4391, 'grad_norm': 8.134387969970703, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 178944, 'step': 3, 'step_time_sec': 4.57, 'avg_step_time_sec': 5.69, 'time_to_completion_sec': 96.71, 'estimated_total_time_sec': 113.77, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 59378.0, 'total_peak_memory_allocated_MB': 38307.87, 'total_peak_memory_reserved_MB': 59378.0, 'step_tokens_per_second': 11598.46, 'avg_tokens_per_second': 11542.88}
+{'loss': 1.3226, 'grad_norm': 7.858397006988525, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 241152, 'step': 4, 'step_time_sec': 6.09, 'avg_step_time_sec': 5.82, 'time_to_completion_sec': 93.18, 'estimated_total_time_sec': 116.48, 'step_peak_memory_allocated_MB': 39875.21, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 10207.14, 'avg_tokens_per_second': 11076.95}
+{'loss': 1.25, 'grad_norm': 4.771732330322266, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 293376, 'step': 5, 'step_time_sec': 4.58, 'avg_step_time_sec': 5.51, 'time_to_completion_sec': 82.69, 'estimated_total_time_sec': 110.26, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 11403.85, 'avg_tokens_per_second': 11144.84}
+{'loss': 1.3064, 'grad_norm': 5.619996070861816, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 346368, 'step': 6, 'step_time_sec': 4.58, 'avg_step_time_sec': 5.33, 'time_to_completion_sec': 74.56, 'estimated_total_time_sec': 106.51, 'step_peak_memory_allocated_MB': 38307.73, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 11580.42, 'avg_tokens_per_second': 11219.69}
+{'loss': 1.3313, 'grad_norm': 4.248163223266602, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 395520, 'step': 7, 'step_time_sec': 3.99, 'avg_step_time_sec': 5.1, 'time_to_completion_sec': 66.34, 'estimated_total_time_sec': 102.06, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 12320.12, 'avg_tokens_per_second': 11363.09}
+{'loss': 1.1924, 'grad_norm': 3.559274673461914, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 464640, 'step': 8, 'step_time_sec': 6.7, 'avg_step_time_sec': 5.33, 'time_to_completion_sec': 63.97, 'estimated_total_time_sec': 106.62, 'step_peak_memory_allocated_MB': 38307.78, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 10315.26, 'avg_tokens_per_second': 11174.94}
+{'loss': 1.2315, 'grad_norm': 3.652539014816284, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 526080, 'step': 9, 'step_time_sec': 6.68, 'avg_step_time_sec': 5.5, 'time_to_completion_sec': 60.5, 'estimated_total_time_sec': 110.0, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 9197.16, 'avg_tokens_per_second': 10874.65}
+{'loss': 1.2283, 'grad_norm': 3.179703950881958, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 580608, 'step': 10, 'step_time_sec': 4.99, 'avg_step_time_sec': 5.44, 'time_to_completion_sec': 54.43, 'estimated_total_time_sec': 108.86, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 10928.59, 'avg_tokens_per_second': 10880.14}
+{'loss': 1.1998, 'grad_norm': 3.2597928047180176, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 638208, 'step': 11, 'step_time_sec': 4.99, 'avg_step_time_sec': 5.4, 'time_to_completion_sec': 48.58, 'estimated_total_time_sec': 107.96, 'step_peak_memory_allocated_MB': 38355.33, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 11533.95, 'avg_tokens_per_second': 10940.63}
+{'loss': 1.1925, 'grad_norm': 3.4332480430603027, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 695808, 'step': 12, 'step_time_sec': 5.47, 'avg_step_time_sec': 5.41, 'time_to_completion_sec': 43.24, 'estimated_total_time_sec': 108.1, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 10520.62, 'avg_tokens_per_second': 10901.95}
+{'loss': 1.1593, 'grad_norm': 3.2433955669403076, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 748032, 'step': 13, 'step_time_sec': 4.63, 'avg_step_time_sec': 5.34, 'time_to_completion_sec': 37.39, 'estimated_total_time_sec': 106.82, 'step_peak_memory_allocated_MB': 38355.33, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 11272.98, 'avg_tokens_per_second': 10928.77}
+{'loss': 1.1761, 'grad_norm': 3.025867462158203, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 801024, 'step': 14, 'step_time_sec': 6.07, 'avg_step_time_sec': 5.4, 'time_to_completion_sec': 32.38, 'estimated_total_time_sec': 107.94, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 8724.41, 'avg_tokens_per_second': 10737.94}
+{'loss': 1.1596, 'grad_norm': 3.086655378341675, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 854016, 'step': 15, 'step_time_sec': 4.98, 'avg_step_time_sec': 5.37, 'time_to_completion_sec': 26.84, 'estimated_total_time_sec': 107.35, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 10639.0, 'avg_tokens_per_second': 10731.38}
+{'loss': 1.171, 'grad_norm': 3.06339430809021, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 905472, 'step': 16, 'step_time_sec': 4.23, 'avg_step_time_sec': 5.29, 'time_to_completion_sec': 21.17, 'estimated_total_time_sec': 105.83, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 12163.45, 'avg_tokens_per_second': 10807.71}
+{'loss': 1.1622, 'grad_norm': 3.0173256397247314, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 964608, 'step': 17, 'step_time_sec': 6.08, 'avg_step_time_sec': 5.34, 'time_to_completion_sec': 16.02, 'estimated_total_time_sec': 106.82, 'step_peak_memory_allocated_MB': 38355.33, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 9723.91, 'avg_tokens_per_second': 10730.58}
+{'loss': 1.1977, 'grad_norm': 3.136826276779175, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1014528, 'step': 18, 'step_time_sec': 3.8, 'avg_step_time_sec': 5.25, 'time_to_completion_sec': 10.5, 'estimated_total_time_sec': 105.01, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 13131.76, 'avg_tokens_per_second': 10832.85}
+{'loss': 1.2126, 'grad_norm': 2.926532745361328, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1068288, 'step': 19, 'step_time_sec': 4.19, 'avg_step_time_sec': 5.19, 'time_to_completion_sec': 5.19, 'estimated_total_time_sec': 103.84, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 12815.75, 'avg_tokens_per_second': 10921.85}
+{'loss': 1.1329, 'grad_norm': 2.8473961353302, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1125888, 'step': 20, 'step_time_sec': 5.12, 'avg_step_time_sec': 5.19, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 103.76, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 11248.78, 'avg_tokens_per_second': 10938.84}
+{'train_runtime': 106.7452, 'train_samples_per_second': 35.974, 'train_steps_per_second': 0.187, 'train_loss': 1.2429506123065948, 'epoch': 0.08, 'num_input_tokens_seen': 1125888, 'step': 20, 'step_time_sec': 5.12, 'avg_step_time_sec': 5.19, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 103.76, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 61730.0, 'total_peak_memory_allocated_MB': 39875.21, 'total_peak_memory_reserved_MB': 61730.0, 'step_tokens_per_second': 11248.78, 'avg_tokens_per_second': 10938.84}
diff --git a/examples/huggingface/llama_B_48_post_init_instance_patching.log b/examples/huggingface/llama_B_48_post_init_instance_patching.log
index 6b0b65012..a3f21f81a 100644
--- a/examples/huggingface/llama_B_48_post_init_instance_patching.log
+++ b/examples/huggingface/llama_B_48_post_init_instance_patching.log
@@ -1,16 +1,21 @@
-***** Pre-init original model *****
-***** Pre-init original model *****
-***** Pre-init original model *****
-***** Pre-init original model *****
-***** Post-init original model *****
-***** Pre-Apply Liger Kernel *****
-***** Post-init original model *****
-***** Pre-Apply Liger Kernel *****
-***** Post-init original model *****
-***** Pre-Apply Liger Kernel *****
-***** Post-init original model *****
-***** Pre-Apply Liger Kernel *****
-***** Post-Apply Liger Kernel *****
-***** Post-Apply Liger Kernel *****
-***** Post-Apply Liger Kernel *****
-***** Post-Apply Liger Kernel *****
+{'loss': 13.3485, 'grad_norm': 535.4854736328125, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 75264}
+{'loss': 13.3631, 'grad_norm': 530.2271118164062, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 132864, 'step': 2, 'step_time_sec': 4.92, 'avg_step_time_sec': 4.92, 'time_to_completion_sec': 88.62, 'estimated_total_time_sec': 98.47, 'step_peak_memory_allocated_MB': 38307.89, 'step_peak_memory_reserved_MB': 54898.0, 'total_peak_memory_allocated_MB': 38307.89, 'total_peak_memory_reserved_MB': 54898.0, 'step_tokens_per_second': 11698.89, 'avg_tokens_per_second': 11698.89}
+{'loss': 10.8637, 'grad_norm': 380.6391296386719, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 185088, 'step': 3, 'step_time_sec': 4.3, 'avg_step_time_sec': 4.61, 'time_to_completion_sec': 78.44, 'estimated_total_time_sec': 92.28, 'step_peak_memory_allocated_MB': 38307.92, 'step_peak_memory_reserved_MB': 55360.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 55360.0, 'step_tokens_per_second': 12131.77, 'avg_tokens_per_second': 11900.82}
+{'loss': 8.3786, 'grad_norm': 47.627689361572266, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 245760, 'step': 4, 'step_time_sec': 5.63, 'avg_step_time_sec': 4.95, 'time_to_completion_sec': 79.25, 'estimated_total_time_sec': 99.07, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 10772.7, 'avg_tokens_per_second': 11473.26}
+{'loss': 8.411, 'grad_norm': 95.4661636352539, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 301056, 'step': 5, 'step_time_sec': 4.32, 'avg_step_time_sec': 4.79, 'time_to_completion_sec': 71.92, 'estimated_total_time_sec': 95.89, 'step_peak_memory_allocated_MB': 38307.92, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 12808.08, 'avg_tokens_per_second': 11773.76}
+{'loss': 8.7209, 'grad_norm': 73.17122650146484, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 354816, 'step': 6, 'step_time_sec': 4.11, 'avg_step_time_sec': 4.66, 'time_to_completion_sec': 65.2, 'estimated_total_time_sec': 93.14, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 13087.07, 'avg_tokens_per_second': 12005.45}
+{'loss': 8.2012, 'grad_norm': 35.51559066772461, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 414720, 'step': 7, 'step_time_sec': 4.18, 'avg_step_time_sec': 4.58, 'time_to_completion_sec': 59.5, 'estimated_total_time_sec': 91.54, 'step_peak_memory_allocated_MB': 38307.91, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 14338.07, 'avg_tokens_per_second': 12360.3}
+{'loss': 7.9819, 'grad_norm': 35.38005828857422, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 466944, 'step': 8, 'step_time_sec': 3.58, 'avg_step_time_sec': 4.43, 'time_to_completion_sec': 53.22, 'estimated_total_time_sec': 88.7, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 14582.17, 'avg_tokens_per_second': 12616.62}
+{'loss': 7.8575, 'grad_norm': 9.446126937866211, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 514560, 'step': 9, 'step_time_sec': 3.47, 'avg_step_time_sec': 4.31, 'time_to_completion_sec': 47.46, 'estimated_total_time_sec': 86.29, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 13723.12, 'avg_tokens_per_second': 12727.86}
+{'loss': 7.7469, 'grad_norm': 7.220946788787842, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 572928, 'step': 10, 'step_time_sec': 5.99, 'avg_step_time_sec': 4.5, 'time_to_completion_sec': 45.0, 'estimated_total_time_sec': 90.01, 'step_peak_memory_allocated_MB': 38307.75, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 9745.02, 'avg_tokens_per_second': 12286.77}
+{'loss': 7.707, 'grad_norm': 5.809014320373535, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 625920, 'step': 11, 'step_time_sec': 3.98, 'avg_step_time_sec': 4.45, 'time_to_completion_sec': 40.04, 'estimated_total_time_sec': 88.97, 'step_peak_memory_allocated_MB': 38307.89, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 13304.21, 'avg_tokens_per_second': 12377.87}
+{'loss': 7.5828, 'grad_norm': 5.793330192565918, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 676608, 'step': 12, 'step_time_sec': 4.11, 'avg_step_time_sec': 4.42, 'time_to_completion_sec': 35.35, 'estimated_total_time_sec': 88.36, 'step_peak_memory_allocated_MB': 38307.78, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 12323.1, 'avg_tokens_per_second': 12373.23}
+{'loss': 7.5844, 'grad_norm': 6.631588459014893, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 725760, 'step': 13, 'step_time_sec': 3.92, 'avg_step_time_sec': 4.38, 'time_to_completion_sec': 30.64, 'estimated_total_time_sec': 87.54, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 12530.66, 'avg_tokens_per_second': 12384.99}
+{'loss': 7.6481, 'grad_norm': 9.671046257019043, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 786432, 'step': 14, 'step_time_sec': 4.48, 'avg_step_time_sec': 4.39, 'time_to_completion_sec': 26.31, 'estimated_total_time_sec': 87.7, 'step_peak_memory_allocated_MB': 38307.91, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 13530.55, 'avg_tokens_per_second': 12475.1}
+{'loss': 7.6636, 'grad_norm': 6.185126781463623, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 844800, 'step': 15, 'step_time_sec': 4.48, 'avg_step_time_sec': 4.39, 'time_to_completion_sec': 21.96, 'estimated_total_time_sec': 87.83, 'step_peak_memory_allocated_MB': 38307.91, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 13038.21, 'avg_tokens_per_second': 12516.1}
+{'loss': 7.5504, 'grad_norm': 8.093387603759766, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 891648, 'step': 16, 'step_time_sec': 4.13, 'avg_step_time_sec': 4.37, 'time_to_completion_sec': 17.5, 'estimated_total_time_sec': 87.49, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 11338.4, 'avg_tokens_per_second': 12441.94}
+{'loss': 7.4777, 'grad_norm': 7.326509952545166, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 956928, 'step': 17, 'step_time_sec': 5.45, 'avg_step_time_sec': 4.44, 'time_to_completion_sec': 13.33, 'estimated_total_time_sec': 88.84, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 11971.5, 'avg_tokens_per_second': 12405.84}
+{'loss': 7.6148, 'grad_norm': 6.828405380249023, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1015296, 'step': 18, 'step_time_sec': 4.6, 'avg_step_time_sec': 4.45, 'time_to_completion_sec': 8.9, 'estimated_total_time_sec': 89.02, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 12696.34, 'avg_tokens_per_second': 12423.49}
+{'loss': 7.5495, 'grad_norm': 7.313730239868164, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1080576, 'step': 19, 'step_time_sec': 4.95, 'avg_step_time_sec': 4.48, 'time_to_completion_sec': 4.48, 'estimated_total_time_sec': 89.58, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 13179.17, 'avg_tokens_per_second': 12469.92}
+{'loss': 7.5819, 'grad_norm': 7.678485870361328, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1138944, 'step': 20, 'step_time_sec': 5.1, 'avg_step_time_sec': 4.51, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 90.23, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 11440.95, 'avg_tokens_per_second': 12408.68}
+{'train_runtime': 95.584, 'train_samples_per_second': 40.174, 'train_steps_per_second': 0.209, 'train_loss': 8.541682934761047, 'epoch': 0.08, 'num_input_tokens_seen': 1138944, 'step': 20, 'step_time_sec': 5.1, 'avg_step_time_sec': 4.51, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 90.23, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 57366.0, 'total_peak_memory_allocated_MB': 38307.92, 'total_peak_memory_reserved_MB': 57366.0, 'step_tokens_per_second': 11440.95, 'avg_tokens_per_second': 12408.68}
diff --git a/examples/huggingface/llama_B_48_post_init_instance_patching_2.log b/examples/huggingface/llama_B_48_post_init_instance_patching_2.log
new file mode 100644
index 000000000..8b237dbf9
--- /dev/null
+++ b/examples/huggingface/llama_B_48_post_init_instance_patching_2.log
@@ -0,0 +1,21 @@
+{'loss': 11.9583, 'grad_norm': 695.6334838867188, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 50688}
+{'loss': 11.9175, 'grad_norm': 682.4520874023438, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 128256, 'step': 2, 'step_time_sec': 6.11, 'avg_step_time_sec': 6.11, 'time_to_completion_sec': 109.94, 'estimated_total_time_sec': 122.16, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 56242.0, 'total_peak_memory_allocated_MB': 38307.87, 'total_peak_memory_reserved_MB': 56242.0, 'step_tokens_per_second': 12699.57, 'avg_tokens_per_second': 12699.57}
+{'loss': 9.8374, 'grad_norm': 272.6592102050781, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 175872, 'step': 3, 'step_time_sec': 3.73, 'avg_step_time_sec': 4.92, 'time_to_completion_sec': 83.64, 'estimated_total_time_sec': 98.4, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 56242.0, 'total_peak_memory_allocated_MB': 38307.87, 'total_peak_memory_reserved_MB': 56242.0, 'step_tokens_per_second': 12757.13, 'avg_tokens_per_second': 12721.4}
+{'loss': 8.667, 'grad_norm': 129.9753875732422, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 235008, 'step': 4, 'step_time_sec': 5.42, 'avg_step_time_sec': 5.09, 'time_to_completion_sec': 81.4, 'estimated_total_time_sec': 101.75, 'step_peak_memory_allocated_MB': 38307.89, 'step_peak_memory_reserved_MB': 57082.0, 'total_peak_memory_allocated_MB': 38307.89, 'total_peak_memory_reserved_MB': 57082.0, 'step_tokens_per_second': 10905.89, 'avg_tokens_per_second': 12076.41}
+{'loss': 8.7041, 'grad_norm': 110.52106475830078, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 291840, 'step': 5, 'step_time_sec': 5.49, 'avg_step_time_sec': 5.19, 'time_to_completion_sec': 77.84, 'estimated_total_time_sec': 103.78, 'step_peak_memory_allocated_MB': 38308.03, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 10345.46, 'avg_tokens_per_second': 11618.29}
+{'loss': 8.083, 'grad_norm': 20.35115623474121, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 350208, 'step': 6, 'step_time_sec': 4.12, 'avg_step_time_sec': 4.98, 'time_to_completion_sec': 69.66, 'estimated_total_time_sec': 99.51, 'step_peak_memory_allocated_MB': 38307.77, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 14164.48, 'avg_tokens_per_second': 12040.05}
+{'loss': 8.312, 'grad_norm': 71.16879272460938, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 409344, 'step': 7, 'step_time_sec': 4.14, 'avg_step_time_sec': 4.84, 'time_to_completion_sec': 62.87, 'estimated_total_time_sec': 96.72, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 14287.55, 'avg_tokens_per_second': 12360.65}
+{'loss': 8.1241, 'grad_norm': 44.23505783081055, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 460800, 'step': 8, 'step_time_sec': 3.94, 'avg_step_time_sec': 4.71, 'time_to_completion_sec': 56.49, 'estimated_total_time_sec': 94.15, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 13071.82, 'avg_tokens_per_second': 12445.6}
+{'loss': 7.8242, 'grad_norm': 14.658903121948242, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 515328, 'step': 9, 'step_time_sec': 3.78, 'avg_step_time_sec': 4.59, 'time_to_completion_sec': 50.51, 'estimated_total_time_sec': 91.83, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 14428.35, 'avg_tokens_per_second': 12649.6}
+{'loss': 8.0857, 'grad_norm': 29.656005859375, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 573696, 'step': 10, 'step_time_sec': 5.46, 'avg_step_time_sec': 4.69, 'time_to_completion_sec': 46.88, 'estimated_total_time_sec': 93.76, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 10691.71, 'avg_tokens_per_second': 12396.27}
+{'loss': 7.8864, 'grad_norm': 10.081645965576172, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 622080, 'step': 11, 'step_time_sec': 3.62, 'avg_step_time_sec': 4.58, 'time_to_completion_sec': 41.23, 'estimated_total_time_sec': 91.62, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 13376.9, 'avg_tokens_per_second': 12473.7}
+{'loss': 8.0753, 'grad_norm': 22.213361740112305, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 703488, 'step': 12, 'step_time_sec': 6.16, 'avg_step_time_sec': 4.72, 'time_to_completion_sec': 37.8, 'estimated_total_time_sec': 94.49, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 13208.49, 'avg_tokens_per_second': 12560.84}
+{'loss': 7.8019, 'grad_norm': 14.009153366088867, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 769536, 'step': 13, 'step_time_sec': 5.54, 'avg_step_time_sec': 4.79, 'time_to_completion_sec': 33.55, 'estimated_total_time_sec': 95.85, 'step_peak_memory_allocated_MB': 38308.03, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 11923.5, 'avg_tokens_per_second': 12499.45}
+{'loss': 7.7286, 'grad_norm': 14.318066596984863, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 825600, 'step': 14, 'step_time_sec': 4.15, 'avg_step_time_sec': 4.74, 'time_to_completion_sec': 28.46, 'estimated_total_time_sec': 94.86, 'step_peak_memory_allocated_MB': 38307.78, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 13508.48, 'avg_tokens_per_second': 12567.37}
+{'loss': 7.8475, 'grad_norm': 12.559386253356934, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 879360, 'step': 15, 'step_time_sec': 4.47, 'avg_step_time_sec': 4.72, 'time_to_completion_sec': 23.62, 'estimated_total_time_sec': 94.47, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 12036.21, 'avg_tokens_per_second': 12531.49}
+{'loss': 7.7093, 'grad_norm': 6.6740899085998535, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 935424, 'step': 16, 'step_time_sec': 4.94, 'avg_step_time_sec': 4.74, 'time_to_completion_sec': 18.95, 'estimated_total_time_sec': 94.76, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 11339.5, 'avg_tokens_per_second': 12448.57}
+{'loss': 7.6226, 'grad_norm': 6.753320217132568, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 998400, 'step': 17, 'step_time_sec': 6.14, 'avg_step_time_sec': 4.83, 'time_to_completion_sec': 14.48, 'estimated_total_time_sec': 96.52, 'step_peak_memory_allocated_MB': 38307.77, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 10253.78, 'avg_tokens_per_second': 12273.99}
+{'loss': 7.6272, 'grad_norm': 6.935516357421875, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1047552, 'step': 18, 'step_time_sec': 4.61, 'avg_step_time_sec': 4.81, 'time_to_completion_sec': 9.63, 'estimated_total_time_sec': 96.27, 'step_peak_memory_allocated_MB': 38307.75, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 10651.27, 'avg_tokens_per_second': 12182.47}
+{'loss': 7.6003, 'grad_norm': 6.853272438049316, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1107456, 'step': 19, 'step_time_sec': 4.14, 'avg_step_time_sec': 4.78, 'time_to_completion_sec': 4.78, 'estimated_total_time_sec': 95.52, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 14468.99, 'avg_tokens_per_second': 12292.59}
+{'loss': 7.7238, 'grad_norm': 6.679880619049072, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1161984, 'step': 20, 'step_time_sec': 5.1, 'avg_step_time_sec': 4.79, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 95.86, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 10694.51, 'avg_tokens_per_second': 12203.12}
+{'train_runtime': 99.2401, 'train_samples_per_second': 38.694, 'train_steps_per_second': 0.202, 'train_loss': 8.456807374954224, 'epoch': 0.08, 'num_input_tokens_seen': 1161984, 'step': 20, 'step_time_sec': 5.1, 'avg_step_time_sec': 4.79, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 95.86, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 60022.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60022.0, 'step_tokens_per_second': 10694.51, 'avg_tokens_per_second': 12203.12}
diff --git a/examples/huggingface/llama_B_48_pre_init_patching.log b/examples/huggingface/llama_B_48_pre_init_patching.log
index 119afdcd7..4438c0be2 100644
--- a/examples/huggingface/llama_B_48_pre_init_patching.log
+++ b/examples/huggingface/llama_B_48_pre_init_patching.log
@@ -1,21 +1,21 @@
-{'loss': 1.424, 'grad_norm': 13.289287567138672, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 60672}
-{'loss': 1.3491, 'grad_norm': 12.90773868560791, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 115968, 'step': 2, 'step_time_sec': 3.92, 'avg_step_time_sec': 3.92, 'time_to_completion_sec': 70.47, 'estimated_total_time_sec': 78.3, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 54698.0, 'total_peak_memory_allocated_MB': 38307.85, 'total_peak_memory_reserved_MB': 54698.0, 'step_tokens_per_second': 14123.62, 'avg_tokens_per_second': 14123.62}
-{'loss': 1.3543, 'grad_norm': 8.59610652923584, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 169728, 'step': 3, 'step_time_sec': 5.11, 'avg_step_time_sec': 4.51, 'time_to_completion_sec': 76.73, 'estimated_total_time_sec': 90.27, 'step_peak_memory_allocated_MB': 38307.99, 'step_peak_memory_reserved_MB': 56882.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 56882.0, 'step_tokens_per_second': 10515.88, 'avg_tokens_per_second': 12080.54}
-{'loss': 1.3071, 'grad_norm': 11.146453857421875, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 216576, 'step': 4, 'step_time_sec': 4.46, 'avg_step_time_sec': 4.5, 'time_to_completion_sec': 71.95, 'estimated_total_time_sec': 89.94, 'step_peak_memory_allocated_MB': 38307.75, 'step_peak_memory_reserved_MB': 58888.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 58888.0, 'step_tokens_per_second': 10494.93, 'avg_tokens_per_second': 11555.91}
-{'loss': 1.1962, 'grad_norm': 4.293575763702393, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 270336, 'step': 5, 'step_time_sec': 4.31, 'avg_step_time_sec': 4.45, 'time_to_completion_sec': 66.74, 'estimated_total_time_sec': 88.99, 'step_peak_memory_allocated_MB': 38307.77, 'step_peak_memory_reserved_MB': 58888.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 58888.0, 'step_tokens_per_second': 12481.65, 'avg_tokens_per_second': 11779.93}
-{'loss': 1.2871, 'grad_norm': 5.704342365264893, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 317184, 'step': 6, 'step_time_sec': 3.56, 'avg_step_time_sec': 4.27, 'time_to_completion_sec': 59.8, 'estimated_total_time_sec': 85.43, 'step_peak_memory_allocated_MB': 38307.77, 'step_peak_memory_reserved_MB': 58888.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 58888.0, 'step_tokens_per_second': 13164.08, 'avg_tokens_per_second': 12010.57}
-{'loss': 1.2825, 'grad_norm': 4.657847881317139, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 380928, 'step': 7, 'step_time_sec': 5.16, 'avg_step_time_sec': 4.42, 'time_to_completion_sec': 57.46, 'estimated_total_time_sec': 88.4, 'step_peak_memory_allocated_MB': 38307.99, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 12348.9, 'avg_tokens_per_second': 12076.43}
-{'loss': 1.2059, 'grad_norm': 3.910782814025879, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 428544, 'step': 8, 'step_time_sec': 3.79, 'avg_step_time_sec': 4.33, 'time_to_completion_sec': 51.97, 'estimated_total_time_sec': 86.61, 'step_peak_memory_allocated_MB': 38307.75, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 12547.12, 'avg_tokens_per_second': 12135.35}
-{'loss': 1.2383, 'grad_norm': 4.36187744140625, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 485376, 'step': 9, 'step_time_sec': 3.82, 'avg_step_time_sec': 4.27, 'time_to_completion_sec': 46.93, 'estimated_total_time_sec': 85.33, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 14879.2, 'avg_tokens_per_second': 12442.39}
-{'loss': 1.1964, 'grad_norm': 3.688460111618042, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 536832, 'step': 10, 'step_time_sec': 3.61, 'avg_step_time_sec': 4.19, 'time_to_completion_sec': 41.94, 'estimated_total_time_sec': 83.88, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 14242.7, 'avg_tokens_per_second': 12614.7}
-{'loss': 1.2785, 'grad_norm': 3.2356770038604736, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 601344, 'step': 11, 'step_time_sec': 4.48, 'avg_step_time_sec': 4.22, 'time_to_completion_sec': 38.01, 'estimated_total_time_sec': 84.46, 'step_peak_memory_allocated_MB': 38307.91, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 14391.72, 'avg_tokens_per_second': 12803.33}
-{'loss': 1.3429, 'grad_norm': 3.291398048400879, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 654336, 'step': 12, 'step_time_sec': 4.18, 'avg_step_time_sec': 4.22, 'time_to_completion_sec': 33.75, 'estimated_total_time_sec': 84.38, 'step_peak_memory_allocated_MB': 38307.91, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 12683.84, 'avg_tokens_per_second': 12792.57}
-{'loss': 1.128, 'grad_norm': 3.3565938472747803, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 713472, 'step': 13, 'step_time_sec': 6.12, 'avg_step_time_sec': 4.38, 'time_to_completion_sec': 30.64, 'estimated_total_time_sec': 87.54, 'step_peak_memory_allocated_MB': 38307.78, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 9669.69, 'avg_tokens_per_second': 12428.95}
-{'loss': 1.2679, 'grad_norm': 3.796682357788086, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 781056, 'step': 14, 'step_time_sec': 6.22, 'avg_step_time_sec': 4.52, 'time_to_completion_sec': 27.11, 'estimated_total_time_sec': 90.37, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 10870.62, 'avg_tokens_per_second': 12264.02}
-{'loss': 1.2586, 'grad_norm': 3.215766191482544, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 832512, 'step': 15, 'step_time_sec': 3.92, 'avg_step_time_sec': 4.48, 'time_to_completion_sec': 22.38, 'estimated_total_time_sec': 89.51, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 13141.74, 'avg_tokens_per_second': 12318.87}
-{'loss': 1.2598, 'grad_norm': 3.1511757373809814, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 893184, 'step': 16, 'step_time_sec': 5.51, 'avg_step_time_sec': 4.54, 'time_to_completion_sec': 18.18, 'estimated_total_time_sec': 90.89, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 11001.5, 'avg_tokens_per_second': 12212.29}
-{'loss': 1.1551, 'grad_norm': 2.9508111476898193, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 966144, 'step': 17, 'step_time_sec': 5.9, 'avg_step_time_sec': 4.63, 'time_to_completion_sec': 13.89, 'estimated_total_time_sec': 92.59, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 12366.52, 'avg_tokens_per_second': 12224.58}
-{'loss': 1.1912, 'grad_norm': 2.9888105392456055, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1037568, 'step': 18, 'step_time_sec': 5.8, 'avg_step_time_sec': 4.7, 'time_to_completion_sec': 9.4, 'estimated_total_time_sec': 93.97, 'step_peak_memory_allocated_MB': 38307.91, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 12312.19, 'avg_tokens_per_second': 12230.94}
-{'loss': 1.2893, 'grad_norm': 3.1170835494995117, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1092864, 'step': 19, 'step_time_sec': 4.13, 'avg_step_time_sec': 4.67, 'time_to_completion_sec': 4.67, 'estimated_total_time_sec': 93.33, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 13395.4, 'avg_tokens_per_second': 12288.17}
-{'loss': 1.222, 'grad_norm': 3.2370665073394775, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1152000, 'step': 20, 'step_time_sec': 6.15, 'avg_step_time_sec': 4.74, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 94.89, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 9620.42, 'avg_tokens_per_second': 12106.26}
-{'train_runtime': 98.8222, 'train_samples_per_second': 38.858, 'train_steps_per_second': 0.202, 'train_loss': 1.2616994500160217, 'epoch': 0.08, 'num_input_tokens_seen': 1152000, 'step': 20, 'step_time_sec': 6.15, 'avg_step_time_sec': 4.74, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 94.89, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 59434.0, 'total_peak_memory_allocated_MB': 38307.99, 'total_peak_memory_reserved_MB': 59434.0, 'step_tokens_per_second': 9620.42, 'avg_tokens_per_second': 12106.26}
+{'loss': 1.437, 'grad_norm': 12.608051300048828, 'learning_rate': 3e-06, 'epoch': 0.0, 'num_input_tokens_seen': 47616}
+{'loss': 1.3617, 'grad_norm': 12.432792663574219, 'learning_rate': 6e-06, 'epoch': 0.01, 'num_input_tokens_seen': 106752, 'step': 2, 'step_time_sec': 4.88, 'avg_step_time_sec': 4.88, 'time_to_completion_sec': 87.76, 'estimated_total_time_sec': 97.51, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 56058.0, 'total_peak_memory_allocated_MB': 38307.87, 'total_peak_memory_reserved_MB': 56058.0, 'step_tokens_per_second': 12129.18, 'avg_tokens_per_second': 12129.18}
+{'loss': 1.2914, 'grad_norm': 8.558877944946289, 'learning_rate': 5.954423259036625e-06, 'epoch': 0.01, 'num_input_tokens_seen': 174336, 'step': 3, 'step_time_sec': 5.8, 'avg_step_time_sec': 5.34, 'time_to_completion_sec': 90.77, 'estimated_total_time_sec': 106.79, 'step_peak_memory_allocated_MB': 38307.75, 'step_peak_memory_reserved_MB': 56058.0, 'total_peak_memory_allocated_MB': 38307.87, 'total_peak_memory_reserved_MB': 56058.0, 'step_tokens_per_second': 11646.39, 'avg_tokens_per_second': 11866.82}
+{'loss': 1.346, 'grad_norm': 8.958698272705078, 'learning_rate': 5.819077862357725e-06, 'epoch': 0.02, 'num_input_tokens_seen': 228864, 'step': 4, 'step_time_sec': 4.25, 'avg_step_time_sec': 4.98, 'time_to_completion_sec': 79.62, 'estimated_total_time_sec': 99.52, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 56058.0, 'total_peak_memory_allocated_MB': 38307.87, 'total_peak_memory_reserved_MB': 56058.0, 'step_tokens_per_second': 12830.81, 'avg_tokens_per_second': 12141.25}
+{'loss': 1.2392, 'grad_norm': 4.964839458465576, 'learning_rate': 5.598076211353317e-06, 'epoch': 0.02, 'num_input_tokens_seen': 274176, 'step': 5, 'step_time_sec': 3.9, 'avg_step_time_sec': 4.71, 'time_to_completion_sec': 70.6, 'estimated_total_time_sec': 94.13, 'step_peak_memory_allocated_MB': 38307.75, 'step_peak_memory_reserved_MB': 56058.0, 'total_peak_memory_allocated_MB': 38307.87, 'total_peak_memory_reserved_MB': 56058.0, 'step_tokens_per_second': 11626.73, 'avg_tokens_per_second': 12034.73}
+{'loss': 1.2877, 'grad_norm': 5.531447887420654, 'learning_rate': 5.298133329356934e-06, 'epoch': 0.02, 'num_input_tokens_seen': 333312, 'step': 6, 'step_time_sec': 5.47, 'avg_step_time_sec': 4.86, 'time_to_completion_sec': 68.02, 'estimated_total_time_sec': 97.17, 'step_peak_memory_allocated_MB': 38308.03, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 10817.49, 'avg_tokens_per_second': 11760.81}
+{'loss': 1.2059, 'grad_norm': 4.537593841552734, 'learning_rate': 4.928362829059618e-06, 'epoch': 0.03, 'num_input_tokens_seen': 387840, 'step': 7, 'step_time_sec': 4.44, 'avg_step_time_sec': 4.79, 'time_to_completion_sec': 62.26, 'estimated_total_time_sec': 95.79, 'step_peak_memory_allocated_MB': 38307.82, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 12271.99, 'avg_tokens_per_second': 11839.85}
+{'loss': 1.2121, 'grad_norm': 3.5951485633850098, 'learning_rate': 4.5e-06, 'epoch': 0.03, 'num_input_tokens_seen': 437760, 'step': 8, 'step_time_sec': 3.94, 'avg_step_time_sec': 4.67, 'time_to_completion_sec': 56.01, 'estimated_total_time_sec': 93.36, 'step_peak_memory_allocated_MB': 38307.84, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 12672.4, 'avg_tokens_per_second': 11940.22}
+{'loss': 1.2958, 'grad_norm': 4.059121608734131, 'learning_rate': 4.0260604299770066e-06, 'epoch': 0.04, 'num_input_tokens_seen': 499968, 'step': 9, 'step_time_sec': 5.0, 'avg_step_time_sec': 4.71, 'time_to_completion_sec': 51.8, 'estimated_total_time_sec': 94.18, 'step_peak_memory_allocated_MB': 38307.98, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 12447.05, 'avg_tokens_per_second': 12007.46}
+{'loss': 1.221, 'grad_norm': 3.738762855529785, 'learning_rate': 3.5209445330007917e-06, 'epoch': 0.04, 'num_input_tokens_seen': 573696, 'step': 10, 'step_time_sec': 6.15, 'avg_step_time_sec': 4.87, 'time_to_completion_sec': 48.7, 'estimated_total_time_sec': 97.39, 'step_peak_memory_allocated_MB': 38307.92, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 11981.58, 'avg_tokens_per_second': 12003.82}
+{'loss': 1.3624, 'grad_norm': 3.72426438331604, 'learning_rate': 3e-06, 'epoch': 0.05, 'num_input_tokens_seen': 628224, 'step': 11, 'step_time_sec': 4.58, 'avg_step_time_sec': 4.84, 'time_to_completion_sec': 43.56, 'estimated_total_time_sec': 96.81, 'step_peak_memory_allocated_MB': 38307.8, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 11907.82, 'avg_tokens_per_second': 11994.74}
+{'loss': 1.2536, 'grad_norm': 3.163553476333618, 'learning_rate': 2.4790554669992093e-06, 'epoch': 0.05, 'num_input_tokens_seen': 692736, 'step': 12, 'step_time_sec': 5.51, 'avg_step_time_sec': 4.9, 'time_to_completion_sec': 39.21, 'estimated_total_time_sec': 98.02, 'step_peak_memory_allocated_MB': 38308.03, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 11715.96, 'avg_tokens_per_second': 11966.27}
+{'loss': 1.2077, 'grad_norm': 3.216358184814453, 'learning_rate': 1.973939570022994e-06, 'epoch': 0.05, 'num_input_tokens_seen': 747264, 'step': 13, 'step_time_sec': 3.81, 'avg_step_time_sec': 4.81, 'time_to_completion_sec': 33.67, 'estimated_total_time_sec': 96.21, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 14296.82, 'avg_tokens_per_second': 12120.25}
+{'loss': 1.2222, 'grad_norm': 3.392892599105835, 'learning_rate': 1.5000000000000007e-06, 'epoch': 0.06, 'num_input_tokens_seen': 818688, 'step': 14, 'step_time_sec': 5.46, 'avg_step_time_sec': 4.86, 'time_to_completion_sec': 29.16, 'estimated_total_time_sec': 97.21, 'step_peak_memory_allocated_MB': 38307.99, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 13071.68, 'avg_tokens_per_second': 12202.52}
+{'loss': 1.2569, 'grad_norm': 3.272480010986328, 'learning_rate': 1.0716371709403819e-06, 'epoch': 0.06, 'num_input_tokens_seen': 867840, 'step': 15, 'step_time_sec': 4.28, 'avg_step_time_sec': 4.82, 'time_to_completion_sec': 24.1, 'estimated_total_time_sec': 96.39, 'step_peak_memory_allocated_MB': 38307.77, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 11477.68, 'avg_tokens_per_second': 12156.52}
+{'loss': 1.2592, 'grad_norm': 2.9783413410186768, 'learning_rate': 7.018666706430663e-07, 'epoch': 0.07, 'num_input_tokens_seen': 937728, 'step': 16, 'step_time_sec': 5.52, 'avg_step_time_sec': 4.87, 'time_to_completion_sec': 19.46, 'estimated_total_time_sec': 97.32, 'step_peak_memory_allocated_MB': 38308.03, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 12668.32, 'avg_tokens_per_second': 12195.2}
+{'loss': 1.2073, 'grad_norm': 3.3658769130706787, 'learning_rate': 4.019237886466839e-07, 'epoch': 0.07, 'num_input_tokens_seen': 994560, 'step': 17, 'step_time_sec': 4.27, 'avg_step_time_sec': 4.83, 'time_to_completion_sec': 14.49, 'estimated_total_time_sec': 96.58, 'step_peak_memory_allocated_MB': 38307.85, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 13300.38, 'avg_tokens_per_second': 12256.32}
+{'loss': 1.2575, 'grad_norm': 3.0143675804138184, 'learning_rate': 1.8092213764227505e-07, 'epoch': 0.07, 'num_input_tokens_seen': 1059840, 'step': 18, 'step_time_sec': 5.01, 'avg_step_time_sec': 4.84, 'time_to_completion_sec': 9.68, 'estimated_total_time_sec': 96.79, 'step_peak_memory_allocated_MB': 38307.98, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 13033.77, 'avg_tokens_per_second': 12303.65}
+{'loss': 1.188, 'grad_norm': 3.1658856868743896, 'learning_rate': 4.557674096337594e-08, 'epoch': 0.08, 'num_input_tokens_seen': 1109760, 'step': 19, 'step_time_sec': 3.82, 'avg_step_time_sec': 4.78, 'time_to_completion_sec': 4.78, 'estimated_total_time_sec': 95.66, 'step_peak_memory_allocated_MB': 38307.87, 'step_peak_memory_reserved_MB': 60588.0, 'total_peak_memory_allocated_MB': 38308.03, 'total_peak_memory_reserved_MB': 60588.0, 'step_tokens_per_second': 13068.22, 'avg_tokens_per_second': 12337.58}
+{'loss': 1.2218, 'grad_norm': 2.7802860736846924, 'learning_rate': 0.0, 'epoch': 0.08, 'num_input_tokens_seen': 1184256, 'step': 20, 'step_time_sec': 6.22, 'avg_step_time_sec': 4.86, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 97.17, 'step_peak_memory_allocated_MB': 38441.34, 'step_peak_memory_reserved_MB': 61934.0, 'total_peak_memory_allocated_MB': 38441.34, 'total_peak_memory_reserved_MB': 61934.0, 'step_tokens_per_second': 11975.59, 'avg_tokens_per_second': 12313.18}
+{'train_runtime': 100.1137, 'train_samples_per_second': 38.356, 'train_steps_per_second': 0.2, 'train_loss': 1.26672260761261, 'epoch': 0.08, 'num_input_tokens_seen': 1184256, 'step': 20, 'step_time_sec': 6.22, 'avg_step_time_sec': 4.86, 'time_to_completion_sec': 0.0, 'estimated_total_time_sec': 97.17, 'step_peak_memory_allocated_MB': 38441.34, 'step_peak_memory_reserved_MB': 61934.0, 'total_peak_memory_allocated_MB': 38441.34, 'total_peak_memory_reserved_MB': 61934.0, 'step_tokens_per_second': 11975.59, 'avg_tokens_per_second': 12313.18}
diff --git a/examples/huggingface/training.py b/examples/huggingface/training.py
index 59e6fcb03..1f5d9be0f 100644
--- a/examples/huggingface/training.py
+++ b/examples/huggingface/training.py
@@ -65,15 +65,33 @@ def train():
     #         torch_dtype=torch.bfloat16,
     #     )
 
+    ## 1. Pre-init patching
+    # _apply_liger_kernel(model_type="llama")
+    # model = transformers.AutoModelForCausalLM.from_pretrained(
+    #     custom_args.model_name,
+    #     trust_remote_code=True,
+    #     use_cache=False,
+    #     torch_dtype=torch.bfloat16,
+    # )
+
+    ## 2. Post-init class-only patching
+    # model = transformers.AutoModelForCausalLM.from_pretrained(
+    #     custom_args.model_name,
+    #     trust_remote_code=True,
+    #     use_cache=False,
+    #     torch_dtype=torch.bfloat16,
+    # )
+    # _apply_liger_kernel(model_type="llama")
+
+    ## 3. Post-init instance patching
     model = transformers.AutoModelForCausalLM.from_pretrained(
         custom_args.model_name,
         trust_remote_code=True,
         use_cache=False,
-        # torch_dtype=torch.bfloat16,
+        torch_dtype=torch.bfloat16,
     )
+    _apply_liger_kernel(model=model) 
 
-    # _apply_liger_kernel(model_type="llama")
-    _apply_liger_kernel(model=model)    
 
     trainer = SFTTrainer(
         model=model,
diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py
index d745ead4a..4d5ac5dcd 100644
--- a/src/liger_kernel/transformers/monkey_patch.py
+++ b/src/liger_kernel/transformers/monkey_patch.py
@@ -78,24 +78,13 @@ def apply_liger_kernel_to_llama(
             # Direct LlamaModel
             base_model = model
        
-        # We need to follow .from_pretrained()
-
-        # # Try reinitializing the base model?
-        # base_model = LlamaModel(config)
-        # setattr(model, "model", base_model)
-
-        base_model.norm = LigerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        torch_dtype = config.torch_dtype
+        base_model.norm = LigerRMSNorm(config.hidden_size, eps=config.rms_norm_eps).to(torch_dtype)
 
         for decoder_layer in base_model.layers:
-            decoder_layer.mlp = LigerSwiGLUMLP(config)
-            decoder_layer.input_layernorm = LigerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-            decoder_layer.post_attention_layernorm = LigerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)  
-
-        # print("***** Calling Model.post_init() *****")
-        base_model.post_init()
-        if base_model != model:
-            model.post_init()
-        # print("***** End calling Model.post_init() *****")
+            decoder_layer.mlp = LigerSwiGLUMLP(config).to(torch_dtype)
+            decoder_layer.input_layernorm = LigerRMSNorm(config.hidden_size, eps=config.rms_norm_eps).to(torch_dtype)
+            decoder_layer.post_attention_layernorm = LigerRMSNorm(config.hidden_size, eps=config.rms_norm_eps).to(torch_dtype)
 
 
 def apply_liger_kernel_to_mistral(