Skip to content

[FEAT][kernels] route LM-head vocab projection through deterministic GEMM #194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
inaniloquentee wants to merge 2 commits into
base: main
Choose a base branch
from
+1 −57
Open

[FEAT][kernels] route LM-head vocab projection through deterministic GEMM #194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6101978
[FEAT][kernels] route LM-head vocab projection through deterministic …
inaniloquentee Jun 26, 2026
b002b29
Merge branch 'main' into feat/deterministic-lm-head-gemm
Flink-ddd Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions rl_engine/executors/deepspeed_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -490,6 +490,7 @@ def _coerce_hidden_tensor(
return hidden
if hasattr(candidate, "hidden_states"):
hidden = _last_hidden_state_tensor(
getattr(candidate, "hidden_states"),
candidate.hidden_states,
expected_hidden_dim=expected_hidden_dim,
Comment on lines 492 to 495

@coderabbitai coderabbitai Bot Jul 20, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win

Remove the duplicate positional argument.

This call passes both getattr(candidate, "hidden_states") and candidate.hidden_states, but _last_hidden_state_tensor accepts only one positional argument. Any object with hidden_states will raise TypeError.

        hidden = _last_hidden_state_tensor(
            getattr(candidate, "hidden_states"),
-           candidate.hidden_states,
            expected_hidden_dim=expected_hidden_dim,
        )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
hidden = _last_hidden_state_tensor(
getattr(candidate, "hidden_states"),
candidate.hidden_states,
expected_hidden_dim=expected_hidden_dim,
hidden = _last_hidden_state_tensor(
getattr(candidate, "hidden_states"),
expected_hidden_dim=expected_hidden_dim,
)
🧰 Tools
🪛 GitHub Actions: CI-Pipeline / 1_linting.txt

[error] 492-492: mypy error: "_last_hidden_state_tensor" gets multiple values for keyword argument "expected_hidden_dim". [misc]

🪛 GitHub Actions: CI-Pipeline / linting

[error] 492-492: mypy: "_last_hidden_state_tensor" gets multiple values for keyword argument "expected_hidden_dim" [misc]

🪛 Ruff (0.15.21)

[warning] 493-493: Do not call getattr with a constant attribute value. It is not any safer than normal property access.

Replace getattr with attribute access

(B009)

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@rl_engine/executors/deepspeed_trainer.py` around lines 492 - 495, Update the
_last_hidden_state_tensor call to pass candidate.hidden_states only once,
removing the duplicate getattr(candidate, "hidden_states") positional argument
while preserving expected_hidden_dim and the remaining call behavior.

)
Expand Down
57 changes: 0 additions & 57 deletions tests/test_linear_logp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,63 +312,6 @@ def test_chunked_linear_logp_backward_matches_autograd_and_layout_invariance(use
assert torch.allclose(grad_bias, canonical_bias_grad, atol=1e-6)


def test_chunked_linear_logp_backward_skips_unused_gradients():
hidden, weight, target, bias = _inputs(2029, device="cpu", bias=True)
grad_out = torch.randn_like(target, dtype=torch.float32)
hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous()
target_1d = target.reshape(-1).contiguous()

full_hidden, full_weight, full_bias = chunked_linear_logp_backward(
grad_out,
hidden_2d,
weight,
target_1d,
bias,
has_bias=True,
lead_shape=target.shape,
hidden_dtype=hidden.dtype,
weight_dtype=weight.dtype,
bias_dtype=bias.dtype,
chunk_elems=weight.size(0) * 3,
)
skipped_hidden, skipped_weight, skipped_bias = chunked_linear_logp_backward(
grad_out,
hidden_2d,
weight,
target_1d,
bias,
has_bias=True,
lead_shape=target.shape,
hidden_dtype=hidden.dtype,
weight_dtype=weight.dtype,
bias_dtype=bias.dtype,
chunk_elems=weight.size(0) * 3,
compute_grad_hidden=False,
)
only_hidden, skipped_weight_2, skipped_bias_2 = chunked_linear_logp_backward(
grad_out,
hidden_2d,
weight,
target_1d,
bias,
has_bias=True,
lead_shape=target.shape,
hidden_dtype=hidden.dtype,
weight_dtype=weight.dtype,
bias_dtype=bias.dtype,
chunk_elems=weight.size(0) * 3,
compute_grad_weight=False,
compute_grad_bias=False,
)

assert skipped_hidden is None
assert torch.allclose(skipped_weight, full_weight, atol=1e-6)
assert torch.allclose(skipped_bias, full_bias, atol=1e-6)
assert skipped_weight_2 is None
assert skipped_bias_2 is None
assert torch.allclose(only_hidden, full_hidden, atol=1e-6)


def test_tied_embedding_lm_head_shared_gradient_is_layout_invariant():
torch.manual_seed(2028)
model = _EmbeddingLMHeadModel(vocab_size=13, hidden_dim=6, bias=False, tie_weights=True)
Expand Down
Loading