diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 45d2baf2b..b73f75b34 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -315,6 +315,8 @@ class _Dumper: model: "torch.nn.Module", name_prefix: str = "param", save: bool = True, + get_grad: Optional[Callable] = None, + step: Optional[int] = None, **kwargs, ) -> None: for param_name, param in model.named_parameters(): @@ -332,6 +334,8 @@ class _Dumper: enable_future_grad=False, value_tag="Dumper.ParamValue", grad_tag="Dumper.ParamGrad", + get_grad=get_grad, + step=step, ) def dump_dict(self, name_prefix, data, save: bool = True, **kwargs): @@ -469,6 +473,8 @@ class _Dumper: value_meta_only_fields: Optional[dict] = None, grad_meta_only_fields: Optional[dict] = None, grafter_extras: Optional[dict] = None, + get_grad: Optional[Callable] = None, + step: Optional[int] = None, ) -> None: self._http_manager # noqa: B018 @@ -499,19 +505,22 @@ class _Dumper: tags=tags, value=value, save=save, + step=step, meta_only_fields={**(value_meta_only_fields or {}), **recompute_meta}, ) - if ( - enable_curr_grad - and isinstance(value, torch.Tensor) - and (g := value.grad) is not None - ): + if enable_curr_grad and isinstance(value, torch.Tensor): + g = get_grad(value) if get_grad is not None else value.grad + else: + g = None + + if g is not None: self._dump_single( tag=grad_tag, tags={**tags, "name": f"grad__{name}"}, value=g, save=save, + step=step, meta_only_fields={**(grad_meta_only_fields or {}), **recompute_meta}, ) diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 70c51c34b..d2767c2a9 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -1127,6 +1127,96 @@ class TestDumpModel: assert all("grad" in f for f in filenames) +class TestDumpModelGradInjection: + def test_get_grad_overrides_param_grad(self, tmp_path): + """dump_model dumps the tensor returned by get_grad instead of param.grad.""" + d = _make_test_dumper( + tmp_path, enable_model_grad=True, enable_model_value=False + ) + model = torch.nn.Linear(4, 2, bias=False) + y = model(torch.ones(1, 4)).sum() + y.backward() + injected = torch.full_like(model.weight, 7.0) + + d.dump_model(model, name_prefix="p", get_grad=lambda param: injected) + + path = _find_dump_file(tmp_path, name="grad__p__weight") + assert torch.equal(_load_dump(path)["value"], injected) + + def test_get_grad_reads_custom_grad_storage(self, tmp_path): + """get_grad lets callers surface grads living outside param.grad (e.g. main_grad).""" + d = _make_test_dumper( + tmp_path, enable_model_grad=True, enable_model_value=False + ) + model = torch.nn.Linear(4, 2, bias=False) + assert model.weight.grad is None + model.weight.main_grad = torch.full_like(model.weight, 3.0) + + d.dump_model( + model, + name_prefix="p", + get_grad=lambda param: getattr(param, "main_grad", None), + ) + + path = _find_dump_file(tmp_path, name="grad__p__weight") + assert torch.equal(_load_dump(path)["value"], model.weight.main_grad) + + def test_get_grad_returning_none_skips_grad_dump(self, tmp_path): + """A get_grad returning None suppresses the grad dump for that param.""" + d = _make_test_dumper( + tmp_path, enable_model_grad=True, enable_model_value=False + ) + model = torch.nn.Linear(4, 2, bias=False) + y = model(torch.ones(1, 4)).sum() + y.backward() + + d.dump_model(model, name_prefix="p", get_grad=lambda param: None) + + assert len(_get_filenames(tmp_path)) == 0 + + def test_default_uses_param_grad_without_main_grad_fallback(self, tmp_path): + """Without get_grad, only param.grad is dumped; a main_grad attribute is ignored.""" + d = _make_test_dumper( + tmp_path, enable_model_grad=True, enable_model_value=False + ) + model = torch.nn.Linear(4, 2, bias=False) + assert model.weight.grad is None + model.weight.main_grad = torch.full_like(model.weight, 3.0) + + d.dump_model(model, name_prefix="p") + + assert len(_get_filenames(tmp_path)) == 0 + + +class TestDumpModelStepOverride: + def test_step_override_in_filenames(self, tmp_path): + """An explicit step overrides the dumper's ambient step for value and grad files.""" + d = _make_test_dumper(tmp_path, enable_model_value=True, enable_model_grad=True) + d._state.step = 3 + model = torch.nn.Linear(4, 2, bias=False) + y = model(torch.ones(1, 4)).sum() + y.backward() + + d.dump_model(model, name_prefix="p", step=7) + + filenames = _get_filenames(tmp_path) + assert len(filenames) == 2 + assert all("step=7" in f for f in filenames) + + def test_step_none_uses_ambient_step(self, tmp_path): + """Without an explicit step, dump_model records the dumper's current step.""" + d = _make_test_dumper( + tmp_path, enable_model_value=True, enable_model_grad=False + ) + d._state.step = 3 + model = torch.nn.Linear(4, 2, bias=False) + + d.dump_model(model, name_prefix="p") + + filenames = _get_filenames(tmp_path) + assert filenames and all("step=3" in f for f in filenames) + + class TestParallelRankInFilename: def test_config_default_false(self): """include_parallel_rank_in_filename defaults to False."""