From cc60c86a378e683a25d40882e12ed1e18d0b7e0f Mon Sep 17 00:00:00 2001 From: "marshall.miao" Date: Wed, 28 Jan 2026 02:49:51 +0000 Subject: [PATCH] relax torch version check to 2.1+ for checkpointers --- bytecheckpoint/checkpointer/ddp_checkpointer.py | 4 ++-- bytecheckpoint/checkpointer/fsdp2_checkpointer.py | 4 ++-- bytecheckpoint/checkpointer/fsdp_checkpointer.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/bytecheckpoint/checkpointer/ddp_checkpointer.py b/bytecheckpoint/checkpointer/ddp_checkpointer.py index 12fabec..54742ce 100644 --- a/bytecheckpoint/checkpointer/ddp_checkpointer.py +++ b/bytecheckpoint/checkpointer/ddp_checkpointer.py @@ -87,8 +87,8 @@ def save( logger.info("Start save function call. ") # Check version. - assert torch.__version__.strip() < "2.6.0" and torch.__version__.strip() >= "2.1.0", ( - "ByteCheckpoint now only support torch version from 2.1 to 2.5" + assert "2.1.0" <= torch.__version__.strip(), ( + "ByteCheckpoint now only support torch version from 2.1 and above" ) # Check role. if role and role not in SUPPORTED_ROLE_SUFFIX_TYPES: diff --git a/bytecheckpoint/checkpointer/fsdp2_checkpointer.py b/bytecheckpoint/checkpointer/fsdp2_checkpointer.py index 7e25cda..0bea947 100644 --- a/bytecheckpoint/checkpointer/fsdp2_checkpointer.py +++ b/bytecheckpoint/checkpointer/fsdp2_checkpointer.py @@ -91,8 +91,8 @@ def save( logger.warning("Unsupported role %s, will use default role instead", role) role = SUPPORTED_ROLE_SUFFIX_TYPES[DEFAULT_STR] # Check version. - assert "2.1.0" <= torch.__version__.strip() and torch.__version__.strip() < "2.6.0", ( - "ByteCheckpoint now only support torch version from 2.1 to 2.5" + assert "2.1.0" <= torch.__version__.strip(), ( + "ByteCheckpoint now only support torch version from 2.1 and above" ) # Check supported components. cls.check_supported_components(checkpoint_state, FSDP_SUPPORTED_TYPES) diff --git a/bytecheckpoint/checkpointer/fsdp_checkpointer.py b/bytecheckpoint/checkpointer/fsdp_checkpointer.py index 4a0f21f..8cd223d 100644 --- a/bytecheckpoint/checkpointer/fsdp_checkpointer.py +++ b/bytecheckpoint/checkpointer/fsdp_checkpointer.py @@ -102,8 +102,8 @@ def save( logger.warning("Unsupported role %s, will use default role instead", role) role = SUPPORTED_ROLE_SUFFIX_TYPES[DEFAULT_STR] # Check version. - assert "2.1.0" <= torch.__version__.strip() and torch.__version__.strip() < "2.6.0", ( - "ByteCheckpoint now only support torch version from 2.1 to 2.5" + assert "2.1.0" <= torch.__version__.strip(), ( + "ByteCheckpoint now only support torch version from 2.1 and above" ) # Check supported components. cls.check_supported_components(checkpoint_state, FSDP_SUPPORTED_TYPES)