feat(serialization): Serialize and deserialize meta tensors

[Eta0]

Serialization and Deserialization of Meta Tensors

This change allows serialization of tensors on the meta device. A brief explanation of meta tensors from the documentation for torch.Tensor.is_meta is:

Meta tensors are like normal tensors, but they carry no data.

As such, when serialized, they have a data_length of 0 and no presence in the file beyond a tensor header, but they are recorded with a real shape. When deserialized, they are yielded as zero-filled tensors of that shape on the requested device.

When verifying a deserialized meta tensor with verify_hash=True (or TensorDeserializer.verify_module()), only the header hashes are checked, as there is no tensor data included in the original CRC32 to compare against. Verifying that a zero-filled tensor is still zero-filled can be achieved manually with not torch.any(tensor), which returns True for tensors of all zeroes and False if any elements are non-zero, and is more efficient than computing a CRC32.

There is room to adjust the implementation to provide an option to deserialize meta tensors back onto the meta device if desired instead of always zero-filling them. Yielding meta tensors would allow more flexibility for the end user.

[dmarx]

lgtm, will extend the test suite to check the new behavior

[Eta0]

This may warrant bumping up to a new tensorizer data version when meta tensors are encountered during serialization, since trying to deserialize a tensor with 0 data_length on an older tensorizer version will likely break down in bizarre ways.

[dmarx]

sharing informative blurb from DM with Eta. out of scope for this PR, but should probably document this more formally in the future (e.g. add to the tensorizer readme)

---

Tensorizer has a code version and a data version

• The code version is the PyPI package version (like 2.7.1)
• The data version (currently 3) is encoded in files for feature compatibility checks
  • The newest version of tensorizer can read any current data version
  • The newest version of tensorizer emits files with the lowest data version applicable based on the features used
  • Older versions of tensorizer can serialize files with data versions they recognize
  • Data version 1 is the baseline
  • Data version 2 adds support for bfloat16 and some other odd datatypes
  • Data version 3 adds encryption
  • Data version 4 (proposed) would add zero-filled tensors

[Eta0]

Tensorizer has a code version and a data version

• The code version is the PyPI package version (like 2.7.1)

• The data version (currently 3) is encoded in files for feature compatibility checks

  • The newest version of tensorizer can read any current data version
  • The newest version of tensorizer emits files with the lowest data version applicable based on the features used
  • Older versions of tensorizer can serialize files with data versions they recognize
  • Data version 1 is the baseline
  • Data version 2 adds support for bfloat16 and some other odd datatypes
  • Data version 3 adds encryption
  • Data version 4 (proposed) would add zero-filled tensors

To add to this quote from me (hiya), this is the current de facto state for compatibility, not de jure; newer versions of tensorizer are not guaranteed to have the same level of backwards compatibility in the serializer or deserializer (though any change that breaks it will probably be a major version bump). This isn't documented because we have never agreed on a rigorous policy for data version compatibility, but what could be documented at the moment is a table of extant tensorizer versions and their possible output data versions + possible input data versions.

[wbrown]

@Eta0 @dmarx Let's add a data version as part of 2.8.0?

[Eta0]

@Eta0 @dmarx Let's add a data version as part of 2.8.0?

@wbrown Merged in the latest changes from main and added a new data version (4).

In data version 4, the previously-unused "file hash" field in the tensorizer file header is now a "feature flags" field, because the old way of indicating encryption metadata wasn't compatible with newer data versions that aren't known up-front, at the beginning of serialization.
Currently, the feature flags field only tracks whether a file is encrypted or not, since that is the only optional file-level attribute. The presence of opaque tensors and meta tensors aren't tracked in the file feature flags, only in the ways they were previously indicated per-tensor.

I also fixed some broken tests and fixed a bug that was preventing meta tensor deserialization from working correctly in plaid_mode. Should be ready to merge now.

[wbrown]

Great job. 👍