From 37029819d59edc17dd8393b1fcd1a419024f0bec Mon Sep 17 00:00:00 2001
From: simpleqt <89645338+simpleqt@users.noreply.github.com>
Date: Tue, 8 Sep 2026 00:41:09 +0800
Subject: [PATCH] docs(runtime): document max_norm and use_graph in
 clip_tensors_by_global_norm

The Args block covered 4 of the 6 parameters; max_norm and use_graph
were missing.

Signed-off-by: simpleqt <89645338+simpleqt@users.noreply.github.com>
---
 deepspeed/runtime/utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py
index 54a8ddc60d26..6a68c877c290 100644
--- a/deepspeed/runtime/utils.py
+++ b/deepspeed/runtime/utils.py
@@ -969,9 +969,12 @@ def clip_tensors_by_global_norm(input_tensors, max_norm=1.0, global_norm=None, m
     """Clip list of tensors by global norm.
     Args:
         input_tensors: List of tensors to be clipped
+        max_norm (float, optional): Max norm used for clipping. Defaults to 1.0
         global_norm (float, optional): Precomputed norm. Defaults to None.
         mpu (optional): model parallelism unit. Defaults to None.
         eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6
+        use_graph (bool, optional): Whether to capture the norm computation in
+            a CUDA graph. Defaults to False
     Returns:
         float: the global norm
     """