Skip to content

v0.3.0: `MemoryMappedTensor`, pickle-free multithreaded serialization and more!
Compare
Choose a tag to compare
@vmoens vmoens released this 31 Jan 14:07
· 1 commit to release/0.3.0 since this release
v0.3.0
09f56f4

In this release we introduce a bunch of exciting features to TensorDict:

  • We deprecate MemmapTensor in favour of MemoryMappedTensor, which is fully backed by torch.Tensor and not numpy anymore. The new API is faster and way more bug-free than it used too. See #541

  • Saving tensordicts on disk can now be done via memmap, memmap_ and memmap_like which all support multithreading. If possible, serialization is pickle free (memmap + json) and torch.save is only used for classes that fail to be serialized with json. Serializing models using tensordict is now 3-10x faster than using torch.save, even for SOTA LLMs such as LLAMA.

  • TensorDict can now carry non tensor data through the NonTensorData class. Assigning non-tensor data can also be done via __setitem__ and they can be retrieved via __getitem__. #601

  • A bunch of new operations have appeared too such as named_apply (apply with key names) or tensordict.auto_batch_size_(), and operations like update can now be achieved for only a subset of keys.

  • Almost all operations in the library are now faster!

  • We are slowing deprecating lazy classes except for LazyStackedTensorDict. Whereas torch.stack used to systematically return a lazy stack, it now returns a dense stack if the set_lazy_legacy(mode) decorator is set to False (which will be the default in the next release). The old behaviour can be set with set_lazy_legacy(True). Lazy stacks can still be obtained using LazyStackedTensorDict.lazy_stack. Appropriate warnings are raised unless you have patched your code accordingly.

What's Changed

  • [Refactor] MemoryMappedTensor by @vmoens in #541
  • [Feature] Multithread memmap by @vmoens in #592
  • [Refactor] Graceful as_tensor by @vmoens in #549
  • [Test] Fix as_tensor test by @vmoens in #551
  • Fix assignment of str-typed value to _device attribute in MemmapTensor by @kurt-stolle in #552
  • [Refactor] Refactor split by @vmoens in #555
  • [Refactor] Refactor implement_for by @vmoens in #556
  • [Feature] Better constructors for MemoryMappedTensors by @vmoens in #557
  • [CI] Fix benchmark on gpu by @vmoens in #560
  • [CI] Add regular benchmarks to CI in PRs without upload by @vmoens in #561
  • [Refactor] Major refactoring of codebase by @vmoens in #559
  • [Benchmark] Benchmark split and chunk by @vmoens in #564
  • [Performance] Faster split, chunk and unbind by @vmoens in #563
  • [Feature] Consolidate functional calls by @vmoens in #565
  • [Refactor] Improve functional call efficiency by @vmoens in #567
  • [Refactor] Do not lock nested tensordict in tensordictparams by @vmoens in #568
  • [Performance] Faster params and buffer registration in TensorDictParams by @vmoens in #569
  • [BugFix] Graceful attribute error exit in TensorDictParams by @vmoens in #571
  • [Refactor] Upgrade pytree import by @vmoens in #573
  • [BugFix] Compatibility with missing _global_parameter_registration_hooks by @vmoens in #574
  • [Feature] Seed workers in TensorDict.map by @vmoens in #562
  • [Performance] Faster update by @vmoens in #572
  • [Performance] Faster to_module by @vmoens in #575
  • [BugFix] _FileHandler for windows by @vmoens in #577
  • [Performance] Faster __init__ by @vmoens in #576
  • [Feature, Test] Add tests for partial update by @vmoens in #578
  • [BugFix] No fallback on TensorDictModule.__getattr__ for private attributes by @vmoens in #579
  • [BugFix] Fix deepcopy of TensorDictParams by @vmoens in #580
  • Add MANIFEST.in by @vmoens in #581
  • [BugFix] Delete parameter/buffer before setting it with regular setattr in to_module by @vmoens in #583
  • [Feature] named_apply and default value in apply by @vmoens in #584
  • [BugFix] Faster empty_like for MemoryMappedTensor by @vmoens in #585
  • [BugFix] Faster empty_like for MemoryMappedTensor (dup) by @vmoens in #586
  • [BugFix] Adapt MemoryMappedTensor for torch < 2.0 by @vmoens in #587
  • [Performance] Make copy_ a no-op if tensors are identical by @vmoens in #588
  • [BugFix] Fix non-blocking arg in copy_ by @vmoens in #590
  • [Feature] Unbind and stack tds in map with chunksize=0 by @vmoens in #589
  • [Performance] Faster dispatch by @vmoens in #487
  • [Feature] Saving metadata of tensorclass by @vmoens in #582
  • [BugFix] Fix osx tests by @vmoens in #591
  • [Feature] Weakref for unlocking tds by @vmoens in #595
  • [BugFix] Fix pickling of weakrefs by @vmoens in #597
  • [Feature] Return early a tensordict created through memmap with multiple threads by @vmoens in #598
  • [CI] Depend on torch nightly for nightly releases by @vmoens in #599
  • [Feature] Storing non-tensor data in tensordicts by @vmoens in #601
  • [Feature, Test] FSDP and DTensors by @vmoens in #600
  • [Minor] Fix type deletion in tensorclass load_memmap by @vmoens in #602
  • [BugFix] Fix ellipsis check by @vmoens in #604
  • [Feature] Best intention stack by @vmoens in #605
  • [Feature] Remove and check for prints in codebase using flake8-print by @vmoens in #603
  • [Doc] Doc revamp by @vmoens in #593
  • [BugFix, Doc] Fix tutorial by @vmoens in #606
  • [BugFix] Fix gh-pages upload by @vmoens in #607
  • [BugFix] Upload content of html directly by @vmoens in #608
  • [Feature] Improve in-place ops for TensorDictParams by @vmoens in #609
  • [BugFix, CI] Fix GPU benchmarks by @vmoens in #611
  • [Feature] inplace to_module by @vmoens in #610
  • [Versioning] Bump v0.3.0 by @vmoens in #613
  • [Feature] Support group specification by @lucifer1004 in #616
  • [Refactor] Remove remaining MemmapTensor references by @vmoens in #617
  • [Tests] Reorder and regroup tests by @vmoens in #614
  • [Performance] Faster set by @vmoens in #619
  • [Performance] Better shared/memmap inheritance and faster exclude by @vmoens in #621
  • [Benchmark] Benchmark select, exclude and empty by @vmoens in #623
  • [Feature] Improve the empty method by @vmoens in #622
  • [BugFix] Fix is_memmap attribute for memmap_like and memmap by @vmoens in #625
  • Bump jinja2 from 3.1.2 to 3.1.3 in /docs by @dependabot in #626
  • [BugFix] Remove shared/memmap inheritance from clone / select / exclude by @vmoens in #624
  • [BugFix] Fix index in list error by @vmoens in #627
  • [Refactor] Make unbind call tensor.unbind by @vmoens in #628
  • [Feature] auto_batch_size_ by @vmoens in #630
  • [BugFix] Fix NonTensorData interaction by @vmoens in #631
  • [Doc] More doc on how to set and get non-tensor data by @vmoens in #632
  • [Feature] _auto_make_functional and _dispatch_td_nn_modules by @vmoens in #633
  • [BugFIx] Fix exclude indent by @vmoens in #637
  • [BugFix] Limit number of threads in workers for .map() by @vmoens in #638
  • [Feature] Robust to lazy_legacy set to false and context managers for reshape ops by @vmoens in #634
  • [Minor] Typo in lazy legacy warnings by @vmoens in #639
  • [Minor] NestedKey typing issues by @vmoens in #640
  • [Feature] nested_keys option in named_apply by @vmoens in #641
  • [Feature] Expose NestedKey in root by @vmoens in #642

New Contributors

Full Changelog: v0.2.1...v0.3.0

Contributors

khwstolle, lucifer1004, and 2 other contributors
Assets 19
Loading