less tight atol

Author: justheuristic

week05_large_models/practice_part2.ipynb

@@ -2468,7 +2468,7 @@
         "        dist.barrier()\n",
         "        if i != rank: continue\n",
         "        print(f\"TParallel outputs ({rank=}):\", tp_output.data, flush=True)\n",
-        "        assert torch.allclose(tp_output, ref_output, atol=1e-6), f\"output mismatch on {rank=}\"\n",
+        "        assert torch.allclose(tp_output, ref_output, atol=1e-5), f\"output mismatch on {rank=}\"\n",
         "\n",
         "    dist.barrier()  # test 2: backward w.r.t. inputs\n",
         "    assert tp_input.grad is None\n",
@@ -2479,7 +2479,7 @@
         "        dist.barrier()\n",
         "        if i != rank: continue\n",
         "        print(f\"TParallel input grad ({rank=}):\", tp_input.grad.data, flush=True)\n",
-        "        assert torch.allclose(tp_input.grad, ref_input_grad, atol=1e-6), f\"input_grad mismatch on {rank=}\"\n"
+        "        assert torch.allclose(tp_input.grad, ref_input_grad, atol=1e-4), f\"input_grad mismatch on {rank=}\"\n"
       ],
       "metadata": {
         "id": "6SrcK3avd3GA"
@@ -3237,4 +3237,4 @@
       }
     }
   ]
-}
+}