fix(nn): prevent neural network layer parameter aliasing - #79
Conversation
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
left a comment
There was a problem hiding this comment.
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_copyreally does produce an independent, still-trainable leaf.contiguous_untrackedallocates new storage, andParam::newappliesvalue.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_datahandles the CPU case directly and otherwise delegates tobackend_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.
Summary / Problem
Param::Cloneintentionally shares anArc<RwLock<Tensor>>so optimizers can keep updating the same parameter handle. TheLinear,Conv2d, andEmbeddinglayers also derivedClone, which made a copied layer silently share trainable parameters with the original. Updating one layer therefore changed the other layer as well.Changes
Param::detached_copy()to copy the current tensor into an independent gradient-accumulating parameter.ClonefromLinear,Conv2d, andEmbedding.deep_clone()to each affected layer for explicit independent copies while preserving layer configuration and optional biases.Tests
cargo test -p oxmera-nn --all-targets(10 passed)cargo clippy -p oxmera-nn --all-targets -- -D warnings(passed)cargo clippy --workspace --all-targets -- -D warnings(passed)cargo check --workspace --all-targets --locked(passed)kernels.ptxdoes not matchkernels.cu; this change does not touch the CUDA backend.Compatibility / Known limitations
Linear,Conv2d, andEmbeddingno longer implement implicitClone. Calldeep_clone()when an independent layer copy is required.Param::Cloneremains a shared-handle clone for optimizer and parameter-list semantics.The local pre-PR image did not include
cargo-denyor 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