From 3420a6f7d165082241db57aba2754c7c07f7cd80 Mon Sep 17 00:00:00 2001
From: urlesistiana <55231606+urlesistiana@users.noreply.github.com>
Date: Wed, 1 Oct 2025 17:55:24 +0800
Subject: [PATCH] check activation_memory_budget value range and accept value 0

---
 library/train_util.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/library/train_util.py b/library/train_util.py
index 8a54cd0c..67df2258 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -5512,11 +5512,18 @@ def prepare_accelerator(args: argparse.Namespace):
     if args.torch_compile:
         dynamo_backend = args.dynamo_backend
 
-    if args.activation_memory_budget:
-        logger.info(
-            f"set torch compile activation memory budget to {args.activation_memory_budget}"
-        )
-        torch._functorch.config.activation_memory_budget = args.activation_memory_budget  # type: ignore
+    if args.activation_memory_budget is not None:  # Note: 0 is a valid value.
+        if 0 <= args.activation_memory_budget <= 1:
+            logger.info(
+                f"set torch compile activation memory budget to {args.activation_memory_budget}"
+            )
+            torch._functorch.config.activation_memory_budget = (  # type: ignore
+                args.activation_memory_budget
+            )
+        else:
+            raise ValueError(
+                "activation_memory_budget must be between 0 and 1 (inclusive)"
+            )
 
     kwargs_handlers = [
         (