Skip to content

feat(core): implement optimized serialization tensor saving#12

Merged
Leahlijuan merged 13 commits intomainfrom
feature/tensor.copy
Jan 30, 2026
Merged

feat(core): implement optimized serialization tensor saving#12
Leahlijuan merged 13 commits intomainfrom
feature/tensor.copy

Conversation

@Leahlijuan
Copy link
Collaborator

  • Implement _save_tensor_optimized in CheckpointSaver to support writing tensors directly to mmap buffers using MLF_TENS format.
  • Update CheckpointLoader to transparently detect and load both the new MLF_TENS format and legacy torch.save format.
  • Enable use_optimized_save by default in NeMo wrapper_util to leverage performance improvements.

- Implement `_save_tensor_optimized` in `CheckpointSaver` to support writing tensors directly to mmap buffers using `MLF_TENS` format.
- Update `CheckpointLoader` to transparently detect and load both the new `MLF_TENS` format and legacy `torch.save` format.
- Enable `use_optimized_save` by default in NeMo `wrapper_util` to leverage performance improvements.
g-husam

This comment was marked as resolved.

@g-husam g-husam changed the title feat(core): implement optimized zero-copy tensor saving feat(core): implement optimized serialization tensor saving Jan 30, 2026
@Leahlijuan Leahlijuan merged commit efd8560 into main Jan 30, 2026
4 checks passed
@Leahlijuan Leahlijuan deleted the feature/tensor.copy branch January 30, 2026 18:08
Leahlijuan added a commit that referenced this pull request Feb 2, 2026
- Implement `_save_tensor_optimized` in `CheckpointSaver` to support
writing tensors directly to mmap buffers using `MLF_TENS` format.
- Update `CheckpointLoader` to transparently detect and load both the
new `MLF_TENS` format and legacy `torch.save` format.
- Enable `use_optimized_save` by default in NeMo `wrapper_util` to
leverage performance improvements.

---------

Co-authored-by: g-husam <husameldawi@google.com>
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.

2 participants