Skip to content

fix(nn): prevent neural network layer parameter aliasing - #79

Merged
vyncint merged 2 commits into
vyncint:mainfrom
mikemikimike:codex/oxmera-60-param-deep-copy
Sep 13, 2026
Merged

vyncint merged 2 commits into
vyncint:mainfrom
mikemikimike:codex/oxmera-60-param-deep-copy

Conversation

@mikemikimike

Copy link
Copy Markdown
Contributor

Summary / Problem

Param::Clone intentionally shares an Arc<RwLock<Tensor>> so optimizers can keep updating the same parameter handle. The Linear, Conv2d, and Embedding layers also derived Clone, which made a copied layer silently share trainable parameters with the original. Updating one layer therefore changed the other layer as well.

Changes

  • Add Param::detached_copy() to copy the current tensor into an independent gradient-accumulating parameter.
  • Remove implicit Clone from Linear, Conv2d, and Embedding.
  • Add deep_clone() to each affected layer for explicit independent copies while preserving layer configuration and optional biases.
  • Add regression coverage for independent weight and bias storage across all three layer types.

Tests

  • cargo test -p oxmera-nn --all-targets (10 passed)
  • cargo clippy -p oxmera-nn --all-targets -- -D warnings (passed)
  • Pre-PR worker: cargo clippy --workspace --all-targets -- -D warnings (passed)
  • Pre-PR worker: cargo check --workspace --all-targets --locked (passed)
  • Pre-PR worker: all non-CUDA workspace tests passed. The full workspace run reaches a pre-existing CUDA PTX fingerprint test failure because the committed kernels.ptx does not match kernels.cu; this change does not touch the CUDA backend.

Compatibility / Known limitations

Linear, Conv2d, and Embedding no longer implement implicit Clone. Call deep_clone() when an independent layer copy is required. Param::Clone remains a shared-handle clone for optimizer and parameter-list semantics.

The local pre-PR image did not include cargo-deny or the repository's pinned research nightly. The repository-wide rustfmt check also reports existing line-ending violations across the upstream baseline.

Issue link

Fixes #60

Signed-off-by: mikemikimike <13286568797@163.com>
…able

Removing `Clone` from `Linear`, `Conv2d` and `Embedding` is a breaking
public API change and the CHANGELOG did not say so. Recorded under
Changed with the reason, the replacement (`deep_clone`,
`Param::detached_copy`), and the note that nothing in the workspace
cloned these layers, so the break is confined to out-of-tree callers.

The new regression test asserts value independence, which is half the
property. A copy that is no longer a gradient-accumulating leaf would
pass every one of those assertions and then silently never learn, so a
second test pins the leaf flag, that the copy accumulates its own
gradient, and that the original does not see it.

Signed-off-by: Vyncint Ng <115854244+vyncint@users.noreply.github.com>

@vyncint vyncint left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

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

Thank you for this. It is a clean fix to a genuinely subtle bug, and the part I most appreciate is that you read Param's sharing as deliberate rather than as the thing to remove. That distinction is the whole issue: the handle must stay shared for optimizers and parameters(), and only the layer copy needed to stop being implicit. You got that boundary exactly right, and documented it on Param itself so the next reader does not have to rediscover it.

I verified the behaviour rather than taking the diff's word for it:

  • Param::detached_copy really does produce an independent, still-trainable leaf. contiguous_untracked allocates new storage, and Param::new applies value.detach().requires_grad_(true), so the copy is a fresh gradient-accumulating leaf rather than a detached constant.
  • The copy stays on its device. contiguous_data handles the CPU case directly and otherwise delegates to backend_for(device)?.contiguous(self), so your doc comment's device claim holds on Metal and CUDA too.
  • Nothing in the workspace, in oxidelake's predict, or in oxmega cloned any of these three layers, so the break has no in-tree or downstream victims.

Full workspace suite here: 167 passed, 0 failed.

I pushed one commit rather than send you round again for it.

The CHANGELOG. Removing Clone from three public types is a breaking API change, and the Unreleased section did not mention it. That matters more than usual right now because this repository has no cargo-semver-checks job yet (that is #59), so the CHANGELOG is currently the only place a break gets declared. I recorded it under Changed with the reason, the replacement, and the note that the blast radius is out-of-tree callers only.

One more assertion. Your regression test pins value independence, which is half the property. The other half is that the copy is still trainable, and a copy that had quietly stopped being a leaf would have passed every assertion in your test and then simply never learned. So there is now a second test asserting the leaf flag, that the copy accumulates its own gradient after a backward pass, and that the original does not see it. It passes against your implementation unchanged, which is the point: it pins behaviour you already got right.

Both additions are formatting-clean and the full required set is green on the merge head.

One small correction to the PR description, purely so it does not mislead anyone reading back. The CUDA PTX fingerprint test is not failing on main; the_shipped_ptx_was_built_from_the_shipped_source passes here, and so does cargo fmt --all --check. Taken together with the line-ending violations your worker also reported, that points at CRLF translation in that checkout rather than anything wrong in the repository. Worth checking core.autocrlf there before it costs you time on a future PR.

Follow-up for me, not for you: LayerNorm and BatchNorm2d also carry parameters and have no deep_clone, and Sequential cannot be copied at all. Your change makes the layer surface consistent in that none of them clone implicitly any more, and I will file the gap separately.

Merging. Nice work.

@vyncint
vyncint merged commit ca56368 into vyncint:main Sep 13, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Clone on Linear, Conv2d and Embedding shares parameters with the original

2 participants