From d5779c7900b63e7df791645e6a5a1302f177ad00 Mon Sep 17 00:00:00 2001
From: Eitan Turok <eitan.turok@gmail.com>
Date: Thu, 26 Sep 2024 22:27:34 +0000
Subject: [PATCH] warn checkpointing does not work

---
 llmfoundry/utils/builders.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/llmfoundry/utils/builders.py b/llmfoundry/utils/builders.py
index 12e734446c..b1e1a303fe 100644
--- a/llmfoundry/utils/builders.py
+++ b/llmfoundry/utils/builders.py
@@ -7,14 +7,9 @@
 import logging
 import os
 import re
+import warnings
 from collections import OrderedDict
-from typing import (
-    Any,
-    ContextManager,
-    Iterable,
-    Optional,
-    Union,
-)
+from typing import Any, ContextManager, Iterable, Optional, Union
 
 import torch
 from composer.core import Algorithm, Callback, Evaluator
@@ -711,6 +706,10 @@ def build_tp_strategies(
     name: str,
     model: ComposerModel,
 ) -> dict[str, ParallelStyle]:
+
+    warnings.warn(
+        'Checkpointing is not currently supported for tensor parallelism due to this pytorch bug: https://github.com/pytorch/pytorch/issues/134095#issuecomment-2345018244'
+    )
     return construct_from_registry(
         name=name,
         registry=registry.tp_strategies,