Skip to content

Fix tensor leakage in discojs - #1237

Merged
JulienVig merged 24 commits into
mainfrom
fix-tensor-leaks
Oct 14, 2026
Merged

JulienVig merged 24 commits into
mainfrom
fix-tensor-leaks

Conversation

@JulienVig

@JulienVig JulienVig commented Sep 25, 2026 •

Copy link
Copy Markdown
Collaborator

Refined tensory memory management (see DISCOJS.md):

  • One owner per tensor: only the owner disposes a tensor. Other code just borrows it and must not dispose it.
  • Function arguments belong to the caller: The function borrows them and must not dispose them.
  • Return values belong to the caller, who must dispose them.
  • Getters (model.weights, .get(i)) return tensors still owned by the object, so the caller must not dispose them.
  • Constructors and wrappers (new WeightsContainer(tensors)) take ownership of the tensors passed in and dispose them in their own dispose().
  • Stored tensors: an object that keeps a tensor after a call (in a field, a map or an aggregator) stores a .clone(), not the argument itself.
  • Explicit handover (trainer.releaseModel()): the caller becomes the owner and the previous owner drops its reference.

Refactored the code to follow these conventions and fix the tensor leaks found along the way:

  • Aggregators: fixed tensor leaks in the byzantine, percentile_clipping and secure_history aggregators. Aggregators are now freed when the client closes, and makePayloads returns copies instead of sharing the original tensors.
  • WeightsContainer: its methods no longer share memory with their inputs. Fixed leaks in equal and reduce, and switched from .dataSync() to the async .data().
  • GPT: the decoded GPT-2 weights, attention mask, token embeddings are now freed. Fixed a leak in modelSync and stopped calls to a model that was already freed.
  • Serialization: models are freed after serialization, and tensors are freed if serialization fails.
  • Training and clients: tensors are freed on error, models that are no longer used are freed, and the trainer, Disco and the federated/decentralized clients were cleaned up.
  • Webapp: inference stops when its component unmounts. The LLM is freed when leaving ChatUI, and the model is freed after creating a new task and in ModelLibrary and TestSteps.
  • CLI, server and onnx-converter: models are freed in the CLI scripts and in the server's model_set.
  • Docs: updated the tensor management rules in DISCOJS.md, and the examples now use trainer.releaseModel.

@JulienVig
JulienVig force-pushed the fix-tensor-leaks branch 3 times, most recently from 10cd762 to 7b599ac Compare October 1, 2026 11:01
@JulienVig
JulienVig marked this pull request as ready for review October 1, 2026 11:02
@JulienVig

JulienVig commented Oct 1, 2026 •

Copy link
Copy Markdown
Collaborator Author

@ale18V @leonard-amsler, please have a look at the updated tensor memory mangement section DISCOJS.md.
TF.js really hasn't been designed to be used outside of scripting so it's really easy to make mistakes (which exactly the opposite of Rust so all these ownsership conventions will be enforced by the compiler once the transition to Burn is done, see #1199)
I will try to merge this PR after you merged yours so that you don't have to solve too many conflicts to resolve. If you create new branches before this one is merged I recomment you base them off this one

@JulienVig
JulienVig added this pull request to stack #1242 October 2, 2026 15:58
@JulienVig
JulienVig force-pushed the fix-tensor-leaks branch 3 times, most recently from 228f1f2 to 47bfc46 Compare October 13, 2026 14:57
@JulienVig
JulienVig force-pushed the fix-tensor-leaks branch 2 times, most recently from b68268f to 28d1999 Compare October 13, 2026 15:32
@JulienVig
JulienVig merged commit 0d8f6fe into main Oct 14, 2026
20 checks passed
@JulienVig
JulienVig deleted the fix-tensor-leaks branch October 14, 2026 11:15
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.

1 participant