Return a copy from OnebitLamb.get_lamb_coeffs

[vineethsaivs]

Problem

OnebitLamb.step() opens by dropping the previous step's stats in place:

# remove the previous stats
del self.lamb_coeffs[:]

and get_lamb_coeffs() handed back that same list object:

def get_lamb_coeffs(self):
    return self.lamb_coeffs

So a caller who reads the coefficients (to log them, or to watch the trust ratios during warmup) is holding the optimizer's own list, and the next step() empties it under them. The snapshot silently becomes [] instead of the values that were read.

Running the two accessors as they exist today, with the one step() statement that touches the list:

OnebitLamb  (deepspeed/runtime/fp16/onebit/lamb.py)
   source            : return self.lamb_coeffs
   before step()     : [tensor(0.5000), tensor(1.5000)]
   after  step()     : []
   caller's snapshot emptied by step(): True

FusedLamb   (deepspeed/ops/lamb/fused_lamb.py)
   source            : return lamb_coeffs
   before step()     : [0.5, 1.5]
   after  step()     : [0.5, 1.5]
   caller's snapshot emptied by step(): False

FusedLamb has the identical del self.lamb_coeffs[:] in its step() and the identical accessor name, and is not affected only because its version builds a new list on the way out. So the two optimizers disagree today about whether the value they hand you survives the next step.

Fix

Return a copy, so both optimizers behave the same way:

return list(self.lamb_coeffs)

One thing I deliberately did not change

FusedLamb.get_lamb_coeffs returns Python floats ([c.item() for c in self.lamb_coeffs]) while this one returns the tensors. Making them match would mean changing the element type this method has always returned, which is a separate call from fixing the aliasing, so I left it alone rather than folding a behaviour change into a bug fix. Happy to align it here or in a follow-up if you would rather the two accessors were identical.

Test

test_onebit_lamb_get_lamb_coeffs_returns_a_copy in tests/unit/runtime/half_precision/onebit/test_onebit.py. It needs no accelerator and no distributed backend: the accessor is pure Python, and since OnebitLamb.__init__ asserts on an initialized backend, the test builds the instance with __new__ and gives it only the attribute the accessor reads. It sits with the other 1-bit Lamb tests rather than in a new file, and it runs in cpu-torch-latest since that job runs all of unit/ and this module has no accelerator-level skip.

Fails before the change (step() emptied the list returned to the caller) and passes after.

Checks

pre-commit run --files deepspeed/runtime/fp16/onebit/lamb.py tests/unit/runtime/half_precision/onebit/test_onebit.py is clean, including yapf, flake8, check-torchdist, check-license and codespell. Commit is signed off.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1f9bf6fcd8

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add the missing Signed-off-by trailer

This reviewed commit is a non-merge commit with one parent, but its commit message has no Signed-off-by: trailer; the repository requires every non-merge commit to include one, so the DCO/custom CI check will reject this change until the commit is amended with --signoff.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.