Repository: coreweave/tensorizer (language: python)

Issue:
Fix S3 uploads and lots of bugs, reorganize files, expose custom S3 config path functionality
# Fixes and Reorganization
This PR contains lots of important bug fixes, small features, and reorganization for some files and method parameters.
Most commits are for a single change, so they can be browsed for more detail.

## Reorganization
- Imports from `tensorizer.serialization` can now be directly imported from `tensorizer`
  - Before: `from tensorizer.serialization import TensorDeserializer`
  - After: `from tensorizer import TensorDeserializer`
- `protobuf.py` is now included in the main `tensorizer` package
- `proto/*`, `tensors/*`, and `CMakeLists.txt` are now included in built source distributions
- Outdated `tests_*.py` files in the root directory were moved to a new `ATTIC` directory for unused code
- `tensorizer/tensors.proto` and `tensorizer/tensors_pb2.py` are now symlinks rather than copies (from `proto/` and `tensors/`, respectively)

## Features
- `stream_io.open_stream` now allows specifying an `s3_config_path` parameter to locate a particular `.s3cfg` file
  - Important when multiple config files are present or when `.s3cfg` is placed in a nonstandard location (e.g. in `/var/run`)
- `stream_io.open_stream` now allows empty-string credentials that do not trigger config parsing, for accessing public buckets that need no credentials
- `TensorDeserializer` instance methods no longer take redundant `device` and `dtype` inputs, since they are already specified in `TensorDeserializer.__init__()`
- `CURLStreamFile` now implements the context management protocol

## Tests
- There are now test cases for `open_stream`'s S3 functionality through `moto` mocked S3 calls

## Fixes
See the commits for the full list. Some highlights:
- S3 upload streams now correctly upload and clean up temporary files even when they aren't explicitly closed
  - They still *should* be closed, but they won't throw unholy errors if someone forgets
- `TensorSerializer` and `TensorDeserializer` have fewer potential errors on `__del__` under other error scenarios
- `TensorType` is now exposed as an importable name from `tensorizer.serialization` (and thus `tensorizer`)
  - It has been required for calls to `TensorSerializer.write_tensor()` but wasn't listed in `__all__`
- Lots of things were broken regarding `device` checks in `TensorDeserializer` and are now fixed (17af2fa42e6503f1741640193d0a16f8ed2feba1)
  - This includes multiple compounding bugs that prevented memory pinning from taking place when `"cuda"` was an inferred device rather than an explicitly specified string
- Method-level `filter_func` checks in `TensorDeserializer` now stack with the `__init__`-level `filter_func` check rather than replacing it
  - This had previously allowed a bug where following a more restrictive `filter_func` by a less restrictive one would lead to a buffer overflow, as the `__init__`-level `filter_func` controlled the buffer size
- `TensorDeserializer`'s `mmap` buffer is no longer eagerly closed when deserializing tensors on the CPU
  - This was resulting in use-after-free segfaults and other nastiness from deallocating memory that was still in use
- `CURLStreamFile` stream file objects are now closed correctly when `curl` exits successfully
- `open_stream` S3 endpoint URLs may now include `https://`, and are rejected if they include `http://` or another scheme
  - Previously, they were required to list no scheme, as `https://` was prepended to them unconditionally
- `open_stream` no longer allows opening local files in text mode
  - This was only allowed for local files anyway, so this usage can and should be replaced with a plain call to `open`
  - See further reasoning in the commit message: f8ec5e6d095852365316214adbec9fd4ad89d045


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
new file mode 100644
index 00000000..cb7a053b
--- /dev/null
+++ b/.github/workflows/build.yml
@@ -0,0 +1,41 @@
+name: Build
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build_dist:
+    name: Build distribution on Ubuntu
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/checkout@v3
+
+      - uses: actions/setup-python@v3
+        with:
+          python-version: '3.8'
+
+      - name: Install build dependencies
+        run: python -m pip install --no-cache-dir -U setuptools build
+
+      - name: Build distribution
+        run: python -m build
+
+      - uses: actions/upload-artifact@v3
+        with:
+          path: ./dist/*
+
+  upload_dist:
+    needs: [ build_dist ]
+    name: Upload distribution to PyPI
+    runs-on: ubuntu-20.04
+    if: github.event_name == 'release' && github.event.action == 'published'
+    steps:
+      - uses: actions/download-artifact@v3
+        with:
+          name: artifact
+          path: dist
+
+      - uses: pypa/gh-action-pypi-publish@release/v1.8.1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 93f49582..01ff6301 100644
--- a/.gitignore
+++ b/.gitignore
@@ -19,3 +19,4 @@ pybuild
 *.pyc
 *.egg-info
 build/
+dist/
diff --git a/ATTIC/README.md b/ATTIC/README.md
new file mode 100644
index 00000000..afed3310
--- /dev/null
+++ b/ATTIC/README.md
@@ -0,0 +1,7 @@
+# Attic (Unused Code)
+
+This directory contains tests and examples for older versions of `tensorizer`
+that have not yet been brought up-to-date.
+
+They do not work with the current version of `tensorizer`,
+but when they do work again, they should be moved out of the attic.
diff --git a/LICENSE b/LICENSE
index 3a851165..4606d594 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,4 +1,4 @@
-Copyright 2022 CoreWeave
+Copyright 2023 CoreWeave
 
 Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 00000000..d31df3be
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,3 @@
+graft tensors
+graft proto
+include CMakeLists.txt
\ No newline at end of file
diff --git a/README.md b/README.md
index 9a5123a8..f8089a39 100644
--- a/README.md
+++ b/README.md
@@ -1,120 +1,371 @@
 # tensorizer
 Module, Model, and Tensor Serialization/Deserialization
 
-## Quick Examples
+## TLDR
+Extremely fast model loads from HTTP/HTTPS and S3 endpoints. GPT-J
+(`20GB`) loads at wire-speed (`~5GB/s`) on a 40GbE network, and is
+only bottlenecked by the Linux kernel TCP stack.
 
-### Transformers
+## Rationale
+CoreWeave and our customers use KNative to deploy models as serverless
+functions. How long a model takes to load is a major factor in the latency
+of KNative scale-up. `tensorizer` is a tool to serialize models and their
+associated tensors into a single file that can be loaded quickly and
+efficiently off an HTTP/HTTPS or S3 endpoint.
 
-```py
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
-from tensorizer import load_model, serialize_model
+By not embedding the model in the container image, we can reduce the
+container image size and the time it takes to load the model. This is
+especially important for models that are large in size, such as
+[EleutherAI/gpt-neox-20B](https://huggingface.co/EleutherAI/gpt-neox-20B)
+that weighs in at `~40GB`.
 
-model_name = "EleutherAI/gpt-neo-125M"
-output_dir = model_name.split("/")[-1]
+This decoupling of the model from the container image also allows us to
+update the model without having to rebuild the container image. This allows
+us to quickly iterate on the model and deploy new versions without having
+to wait for the container image to build or for the container image cache
+to be populated.
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+`tensorizer` has S3 support, so we can store the serialized model in S3
+object storage, and perform streaming loads from S3. This allows us to
+stream the model directly from S3 into the container without having to
+download the model to the container's local filesystem. This also
+pertains to HTTP/HTTPS endpoints, as S3 is just an HTTP/HTTPS endpoint.
 
-# The first model is the CLIP Text Encoder.
-serialize_model(
-    model=model,
-    config=model.config,
-    model_directory=output_dir,
-)
+`tensorizer` also has support for loading models from a local filesystem,
+so you can use it to serialize models locally and load them locally. This
+is extremely fast, as the same principles that make it fast for HTTP/HTTPS
+and S3 endpoints also apply to local filesystems.
+
+## Installation
+
+### From PyPI
+`tensorizer` can be installed from PyPI with `pip`:
+```bash
+python -m pip install tensorizer
+```
+
+### From Source
+You can also install `tensorizer` from source using `pip`.
+
+To clone the repository and install `tensorizer` in
+[editable mode](https://pip.pypa.io/en/stable/topics/local-project-installs/#editable-installs),
+run:
+```bash
+git clone https://github.com/coreweave/tensorizer
+cd tensorizer
+python -m pip install -e .
+```
+
+Or, run the following for `pip` to install `tensorizer`
+[directly from GitHub](https://pip.pypa.io/en/stable/topics/vcs-support/#git):
+```bash
+python -m pip install git+https://github.com/coreweave/tensorizer
+```
+
+## Basic Usage
+Serialization is done with the `TensorSerializer` class. It takes a
+`path_uri` argument that can be a local filesystem path, an HTTP/HTTPS
+endpoint, or an S3 endpoint.
+
+`write_module` is the main method of the `TensorSerializer` class. It
+takes a `torch.nn.Module` and serializes the tensors to the `path_uri`
+endpoint.
+
+The below example serializes the `EleutherAI/gpt-j-6B` model to an S3
+endpoint. It assumes that you have already configured your S3
+credentials in `~/.s3cfg`.
+
+**NOTE:** Loading and serializing `gpt-j-6B` will take a lot of CPU RAM,
+up to `~20GB`. Additionally, when loading `gpt-j-6B` into a GPU, you
+will need about `~16GB` of VRAM. If you don't have that much RAM or VRAM,
+you can use the smaller `gpt-neo-125m` model instead.
 
-# OPTIONAL: You can also save the tokenizer to the output directory.
-tokenizer.save_pretrained(output_dir)
+**NOTE2:** The below examples require the `transformers` and `accelerate`
+libraries. You can install them with `pip`:
+```bash
+python -m pip install transformers accelerate
+```
+
+[serialize.py](examples/serialize.py)
+```python
+from transformers import AutoModelForCausalLM
+from tensorizer import TensorSerializer
+import torch
 
-# To load the models, all we have to do is call load_model().
-model = load_model(
-    path_uri=output_dir,
-    modelclass=AutoModelForCausalLM,
-    configclass=AutoConfig
+model_ref = "EleutherAI/gpt-j-6B"
+# For less intensive requirements, swap above with the line below:
+# model_ref = "EleutherAI/gpt-neo-125m"
+model_name = model_ref.split("/")[-1]
+# Change this to your S3 bucket.
+s3_bucket = "bucket"
+s3_uri = f"s3://{s3_bucket}/{model_name}.tensors"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_ref,
+    revision="float16",
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
 )
 
-# Validate that the model is working. The model is already on the GPU.
-# If you want to use the model on the CPU, you can call model.cpu().
-print(tokenizer.decode(model.generate(
-    tokenizer.encode("I walked my dog", return_tensors="pt").to("cuda"),
-    max_new_tokens=20,
-    pad_token_id=tokenizer.eos_token_id,
-)[0]))
+serializer = TensorSerializer(s3_uri)
+serializer.write_module(model)
+serializer.close()
 ```
 
-### Stable Diffusion
+Conversely, deserialization is done with the `TensorDeserializer` class.
+It takes a `path_uri` argument that can be a local filesystem path, an
+HTTP/HTTPS endpoint, or an S3 endpoint.
+
+`load_into_module` is the main method of the `TensorDeserializer` class.
+It takes a `torch.nn.Module` and loads the tensors from the `path_uri`
+endpoint into the `torch.nn.Module`.
 
-```py
+The below example loads the `EleutherAI/gpt-j-6B` model from an S3
+endpoint.
+
+[deserialize.py](examples/deserialize.py)
+```python
+import torch
 import os
-from diffusers import StableDiffusionPipeline
-from tensorizer import load_model, serialize_model
+import time
+from tensorizer import TensorDeserializer
+from tensorizer.utils import no_init_or_tensor, convert_bytes, get_mem_usage
+from collections import OrderedDict
 
-hf_api_token = os.environ.get("HF_API_TOKEN")
-model_name = "runwayml/stable-diffusion-v1-5"
-model_id = model_name.split("/")[-1]
+# disable missing keys and unexpected key warnings
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"
 
-pipeline = StableDiffusionPipeline.from_pretrained(
-    model_name, use_auth_token=hf_api_token
-)
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 
-# StableDiffusionPipeline is just a collection of models,
-# so we serialize each model in the pipeline individually.
+model_ref = "EleutherAI/gpt-j-6B"
+# To run this at home, swap this with the line below for a smaller example:
+# model_ref = "EleutherAI/gpt-neo-125m"
+model_name = model_ref.split("/")[-1]
+# Change this to your S3 bucket.
+s3_bucket = "bucket"
+s3_uri = f"s3://{s3_bucket}/{model_name}.tensors"
 
-# The first model is the CLIP Text Encoder.
-serialize_model(
-    model=pipeline.text_encoder,
-    config=pipeline.text_encoder.config,
-    model_directory=model_id,
-    model_prefix="encoder",
-)
+config = AutoConfig.from_pretrained(model_ref)
 
-# The second model is the VAE.
-serialize_model(
-    model=pipeline.vae,
-    config=None,
-    model_directory=model_id,
-    model_prefix="vae",
+# This ensures that the model is not initialized.
+model = no_init_or_tensor(
+    lambda: AutoModelForCausalLM.from_pretrained(
+        None, config=config, state_dict=OrderedDict()
+    )
 )
 
-# The third model is the UNet.
-serialize_model(
-    model=pipeline.unet,
-    config=None,
-    model_directory=model_id,
-    model_prefix="unet",
-)
+before_mem = get_mem_usage()
 
-# This is optional, but you can also save the CLIP tokenizer and the Stable
-# Diffusion scheduler to the model directory.
-pipeline.tokenizer.save_pretrained(model_id)
+# Lazy load the tensors from S3 into the model.
+start = time.time()
+deserializer = TensorDeserializer(s3_uri, plaid_mode=True)
+deserializer.load_into_module(model)
+end = time.time()
 
-# We can also save the scheduler.
-pipeline.scheduler.save_config(model_id)
+# Brag about how fast we are.
+total_bytes_str = convert_bytes(deserializer.total_tensor_bytes)
+duration = end - start
+per_second = convert_bytes(deserializer.total_tensor_bytes / duration)
+after_mem = get_mem_usage()
+deserializer.close()
+print(f"Deserialized {total_bytes_str} in {end - start:0.2f}s, {per_second}")
+print(f"Memory usage before: {before_mem}")
+print(f"Memory usage after: {after_mem}")
 
-# To load the models into a blank Stable Diffusion pipeline, we have to import
-# the individual models first so we can initialize the pipeline with them.
+# Tokenize and generate
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained(model_ref)
+input_ids = tokenizer.encode(
+    "¡Hola! Encantado de conocerte. hoy voy a", return_tensors="pt"
+).to("cuda")
 
-from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler
-from transformers import CLIPTextModel, CLIPTextConfig, CLIPTokenizer
+with torch.no_grad():
+    output = model.generate(input_ids, max_new_tokens=50, do_sample=True)
 
-vae = load_model(model_id, AutoencoderKL, None, "vae")
-unet = load_model(model_id, UNet2DConditionModel, None, "unet")
-encoder = load_model(model_id, CLIPTextModel, CLIPTextConfig, "encoder")
+print(f"Output: {tokenizer.decode(output[0], skip_special_tokens=True)}")
+```
 
-pipeline = StableDiffusionPipeline(
-    text_encoder=encoder,
-    vae=vae,
-    unet=unet,
-    scheduler=DDIMScheduler.from_config(model_id),
-    tokenizer=CLIPTokenizer.from_pretrained(model_id),
-    safety_checker=None,
-    feature_extractor=None,
-)
+It should produce output similar to the following, with GPT-J-6B:
+```
+Deserialized model in 6.25 seconds
+Test Output: ¡Hola! Encantado de conocerte. hoy voy a comentar por primera
+vez una teoría de trineo, que quizá te parezca
+algo desconocido, ya que en este mundo han
+llegado a dominar tantos
+```
 
-# Now we can use the pipeline to generate images.
-pipeline.to("cuda")
-pipeline("a photo of an astronaut riding a horse on mars").images[0].save(
-    "image.png"
-)
+More practical examples for the usage of `tensorizer` can be found in
+[examples/hf_serialization.py](examples/hf_serialization.py),
+where `df_main()` serializes models from
+[HuggingFace Diffusers](https://github.com/huggingface/diffusers)
+and `hf_main()` serializes
+[HuggingFace Transformers](https://github.com/huggingface/transformers) models.
+
+## Available Pre-Tensorized Models on the CoreWeave Cloud
+
+The following models are available on the CoreWeave Cloud for free, and can be
+used with the `TensorDeserializer` class. The S3 support defaults to the
+`accel-object.ord1.coreweave.com` endpoint, and the bucket to use `tensorized`.
+
+We name the keys in the S3 bucket after the HuggingFace model identifier, and
+append the `/fp16` suffix for the half-precision version.
+
+For example, the S3 URI for the `EleutherAI/gpt-j-6B` model is:
+`s3://tensorized/EleutherAI/gpt-j-6B/fp16/model.tensors`
+
+The below table shows the available models and their S3 URIs.
+
+### Large Language Models
+
+| Model                                                                     | Precision | S3 URI  |
+|---------------------------------------------------------------------------|-----------|---------|
+| [EleutherAI/gpt-neo-125M](https://huggingface.co/EleutherAI/gpt-neo-125M) | `fp32`    | `s3://tensorized/EleutherAI/gpt-neo-125M/model.tensors` |
+| [EleutherAI/gpt-neo-125M](https://huggingface.co/EleutherAI/gpt-neo-125M) | `fp16`    | `s3://tensorized/EleutherAI/gpt-neo-125M/fp16/model.tensors` |
+| [EleutherAI/gpt-neo-1.3B](https://huggingface.co/EleutherAI/gpt-neo-1.3B) | `fp32`    | `s3://tensorized/EleutherAI/gpt-neo-1.3B/model.tensors` |
+| [EleutherAI/gpt-neo-1.3B](https://huggingface.co/EleutherAI/gpt-neo-1.3B) | `fp16`    | `s3://tensorized/EleutherAI/gpt-neo-1.3B/fp16/model.tensors` |
+| [EleutherAI/gpt-neo-2.7B](https://huggingface.co/EleutherAI/gpt-neo-2.7B) | `fp32`    | `s3://tensorized/EleutherAI/gpt-neo-2.7B/model.tensors` |
+| [EleutherAI/gpt-neo-2.7B](https://huggingface.co/EleutherAI/gpt-neo-2.7B) | `fp16`    | `s3://tensorized/EleutherAI/gpt-neo-2.7B/fp16/model.tensors` |
+| [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B)         | `fp32`    | `s3://tensorized/EleutherAI/gpt-j-6B/model.tensors` |
+| [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B)         | `fp16`    | `s3://tensorized/EleutherAI/gpt-j-6B/fp16/model.tensors` |
+| [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) | `fp32`    | `s3://tensorized/EleutherAI/gpt-neox-20b/model.tensors` |
+| [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) | `fp16`    | `s3://tensorized/EleutherAI/gpt-neox-20b/fp16/model.tensors` |
+| [EleutherAI/pythia-70m](https://huggingface.co/EleutherAI/pythia-70m)     | `fp32`    | `s3://tensorized/EleutherAI/pythia-70m/model.tensors` |
+| [EleutherAI/pythia-70m](https://huggingface.co/EleutherAI/pythia-70m)     | `fp16`    | `s3://tensorized/EleutherAI/pythia-70m/fp16/model.tensors` |
+| [EleutherAI/pythia-1.4b](https://huggingface.co/EleutherAI/pythia-1.4b)   | `fp32`    | `s3://tensorized/EleutherAI/pythia-1.4b/model.tensors` |
+| [EleutherAI/pythia-1.4b](https://huggingface.co/EleutherAI/pythia-1.4b)   | `fp16`    | `s3://tensorized/EleutherAI/pythia-1.4b/fp16/model.tensors` |
+| [EleutherAI/pythia-2.8b](https://huggingface.co/EleutherAI/pythia-2.8b)   | `fp32`    | `s3://tensorized/EleutherAI/pythia-2.8b/model.tensors` |
+| [EleutherAI/pythia-2.8b](https://huggingface.co/EleutherAI/pythia-2.8b)   | `fp16`    | `s3://tensorized/EleutherAI/pythia-2.8b/fp16/model.tensors` |
+| [EleutherAI/pythia-6.9b](https://huggingface.co/EleutherAI/pythia-6.9b)   | `fp32`    | `s3://tensorized/EleutherAI/pythia-6.9b/model.tensors` |
+| [EleutherAI/pythia-6.9b](https://huggingface.co/EleutherAI/pythia-6.9b)   | `fp16`    | `s3://tensorized/EleutherAI/pythia-6.9b/fp16/model.tensors` |
+| [EleutherAI/pythia-12b](https://huggingface.co/EleutherAI/pythia-12b)     | `fp32`    | `s3://tensorized/EleutherAI/pythia-12b/model.tensors` |
+| [EleutherAI/pythia-12b](https://huggingface.co/EleutherAI/pythia-12b)     | `fp16`    | `s3://tensorized/EleutherAI/pythia-12b/fp16/model.tensors` |
+| [EleutherAI/pythia-70m-deduped](https://huggingface.co/EleutherAI/pythia-70m-deduped) | `fp32`    | `s3://tensorized/EleutherAI/pythia-70m-deduped/model.tensors` |
+| [EleutherAI/pythia-70m-deduped](https://huggingface.co/EleutherAI/pythia-70m-deduped) | `fp16`    | `s3://tensorized/EleutherAI/pythia-70m-deduped/fp16/model.tensors` |
+| [EleutherAI/pythia-1.4b-deduped](https://huggingface.co/EleutherAI/pythia-1.4b-deduped) | `fp32`    | `s3://tensorized/EleutherAI/pythia-1.4b-deduped/model.tensors` |
+| [EleutherAI/pythia-1.4b-deduped](https://huggingface.co/EleutherAI/pythia-1.4b-deduped) | `fp16`    | `s3://tensorized/EleutherAI/pythia-1.4b-deduped/fp16/model.tensors` |
+| [EleutherAI/pythia-2.8b-deduped](https://huggingface.co/EleutherAI/pythia-2.8b-deduped) | `fp32`    | `s3://tensorized/EleutherAI/pythia-2.8b-deduped/model.tensors` |
+| [EleutherAI/pythia-2.8b-deduped](https://huggingface.co/EleutherAI/pythia-2.8b-deduped) | `fp16`    | `s3://tensorized/EleutherAI/pythia-2.8b-deduped/fp16/model.tensors` |
+| [EleutherAI/pythia-6.9b-deduped](https://huggingface.co/EleutherAI/pythia-6.9b-deduped) | `fp32`    | `s3://tensorized/EleutherAI/pythia-6.9b-deduped/model.tensors` |
+| [EleutherAI/pythia-6.9b-deduped](https://huggingface.co/EleutherAI/pythia-6.9b-deduped) | `fp16`    | `s3://tensorized/EleutherAI/pythia-6.9b-deduped/fp16/model.tensors` |
+| [EleutherAI/pythia-12b-deduped](https://huggingface.co/EleutherAI/pythia-12b-deduped) | `fp32`    | `s3://tensorized/EleutherAI/pythia-12b-deduped/model.tensors` |
+| [EleutherAI/pythia-12b-deduped](https://huggingface.co/EleutherAI/pythia-12b-deduped) | `fp16`    | `s3://tensorized/EleutherAI/pythia-12b-deduped/fp16/model.tensors` |
+| [KoboldAI/fairseq-dense-125M](https://huggingface.co/KoboldAI/fairseq-dense-125M) | `fp32`    | `s3://tensorized/KoboldAI/fairseq-dense-125M/model.tensors` |
+| [KoboldAI/fairseq-dense-125M](https://huggingface.co/KoboldAI/fairseq-dense-125M) | `fp16`    | `s3://tensorized/KoboldAI/fairseq-dense-125M/fp16/model.tensors` |
+| [KoboldAI/fairseq-dense-355M](https://huggingface.co/KoboldAI/fairseq-dense-355M) | `fp32`    | `s3://tensorized/KoboldAI/fairseq-dense-355M/model.tensors` |
+| [KoboldAI/fairseq-dense-355M](https://huggingface.co/KoboldAI/fairseq-dense-355M) | `fp16`    | `s3://tensorized/KoboldAI/fairseq-dense-355M/fp16/model.tensors` |
+| [KoboldAI/fairseq-dense-2.7B](https://huggingface.co/KoboldAI/fairseq-dense-2.7B) | `fp32`    | `s3://tensorized/KoboldAI/fairseq-dense-2.7B/model.tensors` |
+| [KoboldAI/fairseq-dense-2.7B](https://huggingface.co/KoboldAI/fairseq-dense-2.7B) | `fp16`    | `s3://tensorized/KoboldAI/fairseq-dense-2.7B/fp16/model.tensors` |
+| [KoboldAI/fairseq-dense-6.7B](https://huggingface.co/KoboldAI/fairseq-dense-6.7B) | `fp32`    | `s3://tensorized/KoboldAI/fairseq-dense-6.7B/model.tensors` |
+| [KoboldAI/fairseq-dense-6.7B](https://huggingface.co/KoboldAI/fairseq-dense-6.7B) | `fp16`    | `s3://tensorized/KoboldAI/fairseq-dense-6.7B/fp16/model.tensors` |
+| [KoboldAI/fairseq-dense-13B](https://huggingface.co/KoboldAI/fairseq-dense-13B) | `fp32`    | `s3://tensorized/KoboldAI/fairseq-dense-13B/model.tensors` |
+| [KoboldAI/fairseq-dense-13B](https://huggingface.co/KoboldAI/fairseq-dense-13B) | `fp16`    | `s3://tensorized/KoboldAI/fairseq-dense-13B/fp16/model.tensors` |
+| [Salesforce/codegen-350M-mono](https://huggingface.co/Salesforce/codegen-350M-mono) | `fp32`    | `s3://tensorized/Salesforce/codegen-350M-mono/model.tensors` |
+| [Salesforce/codegen-350M-mono](https://huggingface.co/Salesforce/codegen-350M-mono) | `fp16`    | `s3://tensorized/Salesforce/codegen-350M-mono/fp16/model.tensors` |
+| [Salesforce/codegen-350M-multi](https://huggingface.co/Salesforce/codegen-350M-multi) | `fp32`    | `s3://tensorized/Salesforce/codegen-350M-multi/model.tensors` |
+| [Salesforce/codegen-350M-multi](https://huggingface.co/Salesforce/codegen-350M-multi) | `fp16`    | `s3://tensorized/Salesforce/codegen-350M-multi/fp16/model.tensors` |
+| [Salesforce/codegen-2B-multi](https://huggingface.co/Salesforce/codegen-2B-multi) | `fp32`    | `s3://tensorized/Salesforce/codegen-2B-multi/model.tensors` |
+| [Salesforce/codegen-2B-multi](https://huggingface.co/Salesforce/codegen-2B-multi) | `fp16`    | `s3://tensorized/Salesforce/codegen-2B-multi/fp16/model.tensors` |
+| [Salesforce/codgen-6B-mono](https://huggingface.co/Salesforce/codgen-6B-mono) | `fp32`    | `s3://tensorized/Salesforce/codgen-6B-mono/model.tensors` |
+| [Salesforce/codgen-6B-mono](https://huggingface.co/Salesforce/codgen-6B-mono) | `fp16`    | `s3://tensorized/Salesforce/codgen-6B-mono/fp16/model.tensors` |
+| [Salesforce/codgen-6B-multi](https://huggingface.co/Salesforce/codgen-6B-multi) | `fp32`    | `s3://tensorized/Salesforce/codgen-6B-multi/model.tensors` |
+| [Salesforce/codgen-6B-multi](https://huggingface.co/Salesforce/codgen-6B-multi) | `fp16`    | `s3://tensorized/Salesforce/codgen-6B-multi/fp16/model.tensors` |
+| [Salesforce/codegen-16B-mono](https://huggingface.co/Salesforce/codegen-16B-mono) | `fp32`    | `s3://tensorized/Salesforce/codegen-16B-mono/model.tensors` |
+| [Salesforce/codegen-16B-mono](https://huggingface.co/Salesforce/codegen-16B-mono) | `fp16`    | `s3://tensorized/Salesforce/codegen-16B-mono/fp16/model.tensors` |
+| [Salesforce/codegen-16B-multi](https://huggingface.co/Salesforce/codegen-16B-multi) | `fp32`    | `s3://tensorized/Salesforce/codegen-16B-multi/model.tensors` |
+| [Salesforce/codegen-16B-multi](https://huggingface.co/Salesforce/codegen-16B-multi) | `fp16`    | `s3://tensorized/Salesforce/codegen-16B-multi/fp16/model.tensors` |
+
+## S3 Usage Notes
+`tensorizer` uses the `boto3` library to interact with S3. The easiest way
+to use `tensorizer` with S3 is to configure your S3 credentials in
+`~/.s3cfg`.
+
+If you don't want to use `~/.s3cfg`, or wish to use a `.s3cfg` config file
+saved at a nonstandard location (e.g. under `/var/run`), you can also specify
+your S3 credentials using the `tensorizer.stream_io.open_stream()` function,
+and then pass that into the `TensorSerializer` or `TensorDeserializer`
+constructor.
+
+The `stream_io.open_stream()` function takes a `path_uri` argument, which can
+be an `s3://` URI, and accepts the following keyword arguments:
+* `s3_access_key_id`: S3 access key ID
+* `s3_secret_access_key`: S3 secret access key
+* `s3_endpoint`: S3 endpoint
+
+*Or,*
+
+* `s3_config_path`: Alternative filesystem path to a `.s3cfg` config file
+
+For example:
+```python
+TensorSerializer(
+    open_stream(s3_uri,
+                "wb",
+                s3_access_key_id=ACCESS_KEY,
+                s3_secret_access_key=SECRET_KEY,
+                s3_endpoint="object.ord1.coreweave.com"))
+```
+
+and...
+
+```python
+TensorDeserializer(
+    open_stream(s3_uri,
+                "rb",
+                s3_access_key_id=ACCESS_KEY,
+                s3_secret_access_key=SECRET_KEY,
+                s3_endpoint="object.ord1.coreweave.com"))
 ```
 
-More practical examples for usage of the Tensorizer can be found inside of [tensorizer.py](tensorizer.py), where `df_main()` serializes models from [HuggingFace Diffusers](https://github.com/huggingface/diffusers) and `hf_main()` serializes [HuggingFace Transformers](https://github.com/huggingface/transformers) models.
+**NOTE:** For faster object downloads in the CoreWeave Cloud, you can use
+the `accel-object.ord1.coreweave.com` endpoint. This endpoint is optimized
+for object downloads, and will be faster than the `object.ord1.coreweave.com`
+endpoint once the object is cached.
+
+**NOTE2:** The cache above does not get invalidated when the object is updated
+in S3. If you update an object in S3, you will need to wait for the cache to
+expire before you can download the updated object. This takes 24 hours since
+the last download.
+
+For this reason, it is recommended to use an unique S3 key for each version
+of a model if you use the `accel-object.ord1.coreweave.com` endpoint.
+
+## Additional Features
+`tensorizer` has a few additional features that make it more useful than
+just a serialization/deserialization tool.
+
+### Plaid Mode
+`tensorizer` has a `plaid_mode` argument that can be passed to the
+`TensorDeserializer` class. When `plaid_mode` is `True`, `tensorizer`
+will load the tensors extremely fast. This is done by loading the tensors
+into a `torch.nn.Module` that is not initialized, by overriding the
+`__init__` method of the `torch.nn.Module` to do nothing.
+
+The tensors are them loaded into a buffer, and the buffer is zero-copied
+into the uninitialized `torch.nn.Module`. This is unsafe, and should only
+be used in inference cases where the model is not being trained.
+
+### `state_dict` Support
+The `TensorDeserializer` object can be used as-is as a `state_dict` for
+`torch.nn.Module.load_state_dict`. This is useful for loading the tensors
+into a `torch.nn.Module` that is already initialized, or for inspection.
+
+Keep in mind that `load_state_dict` is not a fast operation, and will
+likely be much slower than `load_into_module`.
+
+The `state_dict` can also be used to initialize a HuggingFace Transformers
+AutoModel. But HuggingFace Transformers performs three or more copies of
+the data, so memory use will explode.
+
+## Running Tests
+`tensorizer` uses `unittest` for testing.
+The tests have their own set of dependencies, which can be installed with
+`pip install -r tests/requirements.txt`.
+
+Some tests require a GPU, and will be skipped if no GPU is available.
+To run the tests, run the following in the root of the repository:
+
+```bash
+python -m pip install -e .
+python -m pip install -r tests/requirements.txt
+python -m unittest discover tests/ --verbose
+```
diff --git a/examples/__init__.py b/examples/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/deserialize.py b/examples/deserialize.py
new file mode 100644
index 00000000..5c8c7f3d
--- /dev/null
+++ b/examples/deserialize.py
@@ -0,0 +1,58 @@
+import torch
+import os
+import time
+from tensorizer import TensorDeserializer
+from tensorizer.utils import no_init_or_tensor, convert_bytes, get_mem_usage
+from collections import OrderedDict
+
+# disable missing keys and unexpected key warnings
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+
+model_ref = "EleutherAI/gpt-j-6B"
+# To run this at home, swap this with the line below for a smaller example:
+# model_ref = "EleutherAI/gpt-neo-125m"
+model_name = model_ref.split("/")[-1]
+# Change this to your S3 bucket.
+s3_bucket = "bucket"
+s3_uri = f"s3://{s3_bucket}/{model_name}.tensors"
+
+config = AutoConfig.from_pretrained(model_ref)
+
+# This ensures that the model is not initialized.
+model = no_init_or_tensor(
+    lambda: AutoModelForCausalLM.from_pretrained(
+        None, config=config, state_dict=OrderedDict()
+    )
+)
+
+before_mem = get_mem_usage()
+
+# Lazy load the tensors from S3 into the model.
+start = time.time()
+deserializer = TensorDeserializer(s3_uri, plaid_mode=True)
+deserializer.load_into_module(model)
+end = time.time()
+
+# Brag about how fast we are.
+total_bytes_str = convert_bytes(deserializer.total_tensor_bytes)
+duration = end - start
+per_second = convert_bytes(deserializer.total_tensor_bytes / duration)
+after_mem = get_mem_usage()
+deserializer.close()
+print(f"Deserialized {total_bytes_str} in {end - start:0.2f}s, {per_second}")
+print(f"Memory usage before: {before_mem}")
+print(f"Memory usage after: {after_mem}")
+
+# Tokenize and generate
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained(model_ref)
+input_ids = tokenizer.encode(
+    "¡Hola! Encantado de conocerte. hoy voy a", return_tensors="pt"
+).to("cuda")
+
+with torch.no_grad():
+    output = model.generate(input_ids, max_new_tokens=50, do_sample=True)
+
+print(f"Output: {tokenizer.decode(output[0], skip_special_tokens=True)}")
diff --git a/examples/hf_serialization.py b/examples/hf_serialization.py
new file mode 100644
index 00000000..43722cb6
--- /dev/null
+++ b/examples/hf_serialization.py
@@ -0,0 +1,317 @@
+import argparse
+import json
+import os
+from collections import OrderedDict
+from typing import Optional, Union
+
+import tensorizer.utils as utils
+from tensorizer import TensorSerializer, TensorDeserializer
+import tensorizer.stream_io as stream_io
+import logging
+import time
+import torch
+import tempfile
+
+os.environ[
+    "TRANSFORMERS_VERBOSITY"
+] = "error"  # disable missing keys and unexpected key warnings
+
+from transformers import (
+    AutoConfig,
+    CLIPTextModel,
+    CLIPTokenizer,
+    CLIPTextConfig,
+)
+from transformers.modeling_utils import PreTrainedModel
+
+from diffusers import (
+    AutoencoderKL,
+    UNet2DConditionModel,
+    StableDiffusionPipeline,
+    LMSDiscreteScheduler,
+)
+from diffusers.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin
+
+# Setup logger
+logger = logging.getLogger(__name__)
+logger.setLevel(level=logging.INFO)
+fh = logging.StreamHandler()
+fh_formatter = logging.Formatter(
+    "%(asctime)s %(levelname)s %(filename)s(%(process)d) - %(message)s"
+)
+fh.setFormatter(fh_formatter)
+logger.addHandler(fh)
+
+
+def serialize_model(
+    model: torch.nn.Module,
+    config: Optional[Union[ConfigMixin, AutoConfig, dict]],
+    model_directory: str,
+    model_prefix: str = "model",
+):
+    """
+    Remove the tensors from a PyTorch model, convert them to NumPy
+    arrays and serialize them to GooseTensor format. The stripped
+    model is also serialized to pytorch format.
+
+    Args:
+        model: The model to serialize.
+        config: The model's configuration. This is optional and only
+            required for HuggingFace Transformers models. Diffusers
+            models do not require this.
+        model_directory: The directory to save the serialized model to.
+        model_prefix: The prefix to use for the serialized model files. This
+            is purely optional and it allows for multiple models to be
+            serialized to the same directory. A good example are Stable
+            Diffusion models. Default is "model".
+    """
+
+    os.makedirs(model_directory, exist_ok=True)
+    dir_prefix = f"{model_directory}/{model_prefix}"
+
+    if config is None:
+        config = model
+    if config is not None:
+        if hasattr(config, "to_json_file"):
+            config.to_json_file(f"{dir_prefix}-config.json")
+        if isinstance(config, dict):
+            open(f"{dir_prefix}-config.json", "w").write(
+                json.dumps(config, indent=2)
+            )
+
+    ts = TensorSerializer(f"{dir_prefix}.tensors")
+    ts.write_module(model)
+    ts.close()
+
+
+def load_model(
+    path_uri: str,
+    modelclass: Union[PreTrainedModel, ModelMixin, ConfigMixin] = None,
+    configclass: Optional[Union[ConfigMixin, AutoConfig]] = None,
+    model_prefix: str = "model",
+    device: torch.device = utils.get_device(),
+    dtype: str = None,
+) -> torch.nn.Module:
+    """
+    Given a path prefix, load the model with a custom extension
+
+    Args:
+        path_uri: path to the model. Can be a local path or a URI
+        modelclass: The model class to load the tensors into.
+        configclass: The config class to load the model config into. This must be
+            set if you are loading a model from HuggingFace Transformers.
+        model_prefix: The prefix to use to distinguish between multiple serialized
+            models. The default is "model".
+        dtype: The dtype to load the tensors into. If None, the dtype is inferred from
+            the model.
+    """
+
+    if model_prefix is None:
+        model_prefix = "model"
+
+    begin_load = time.time()
+    ram_usage = utils.get_mem_usage()
+
+    config_uri = f"{path_uri}/{model_prefix}-config.json"
+    tensors_uri = f"{path_uri}/{model_prefix}.tensors"
+    tensor_stream = stream_io.open_stream(tensors_uri)
+
+    logger.info(f"Loading {tensors_uri}, {ram_usage}")
+
+    tensor_deserializer = TensorDeserializer(
+        tensor_stream, device=device, dtype=dtype, lazy_load=True
+    )
+
+    if configclass is not None:
+        try:
+            with tempfile.TemporaryDirectory() as dir:
+                open(os.path.join(dir, "config.json"), "w").write(
+                    stream_io.open_stream(config_uri).read().decode("utf-8")
+                )
+                config = configclass.from_pretrained(dir)
+                config.gradient_checkpointing = True
+        except ValueError:
+            config = configclass.from_pretrained(config_uri)
+        model = utils.no_init_or_tensor(
+            lambda: modelclass.from_pretrained(
+                None, config=config, state_dict=OrderedDict()
+            )
+        )
+    else:
+        try:
+            config = json.loads(
+                stream_io.open_stream(config_uri).read().decode("utf-8")
+            )
+        except ValueError:
+            with open(config_uri, "r") as f:
+                config = json.load(f)
+        model = utils.no_init_or_tensor(lambda: modelclass(**config))
+
+    tensor_deserializer.load_into_module(model)
+
+    tensor_load_s = time.time() - begin_load
+    rate_str = utils.convert_bytes(
+        tensor_deserializer.total_bytes_read / tensor_load_s
+    )
+    tensors_sz = utils.convert_bytes(tensor_deserializer.total_bytes_read)
+    logger.info(
+        f"Model tensors loaded in {tensor_load_s:0.2f}s, read "
+        + f"{tensors_sz} @ {rate_str}/s, {utils.get_mem_usage()}"
+    )
+
+    return model
+
+
+def df_main(args: argparse.Namespace) -> None:
+    output_prefix = args.output_prefix
+    print("MODEL PATH:", args.input_directory)
+    print("OUTPUT PREFIX:", output_prefix)
+
+    hf_api_token = os.environ.get("HF_API_TOKEN")
+
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        args.input_directory, use_auth_token=hf_api_token
+    )
+
+    logger.info("Serializing model")
+    logger.info("GPU: " + utils.get_gpu_name())
+    logger.info("PYTHON USED RAM: " + utils.get_mem_usage())
+
+    serialize_model(
+        pipeline.text_encoder.eval(),
+        pipeline.text_encoder.config,
+        output_prefix,
+        "encoder",
+    )
+    serialize_model(pipeline.vae.eval(), None, output_prefix, "vae")
+    serialize_model(pipeline.unet.eval(), None, output_prefix, "unet")
+
+    pipeline.tokenizer.save_pretrained(output_prefix)
+
+    if args.validate:
+        device = utils.get_device()
+
+        logger.info("Validating serialization")
+        vae = load_model(output_prefix, AutoencoderKL, None, "vae", device)
+        unet = load_model(output_prefix, UNet2DConditionModel, None, "unet")
+        encoder = load_model(
+            output_prefix, CLIPTextModel, CLIPTextConfig, "encoder", device
+        )
+
+        pipeline = StableDiffusionPipeline(
+            text_encoder=encoder,
+            vae=vae,
+            unet=unet,
+            tokenizer=CLIPTokenizer.from_pretrained(
+                args.input_directory, subfolder="tokenizer"
+            ),
+            scheduler=LMSDiscreteScheduler(
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                beta_start=0.00085,
+                num_train_timesteps=1000,
+                trained_betas=None,
+            ),
+            safety_checker=None,
+            feature_extractor=None,
+        ).to(device)
+
+        prompt = "a photo of an astronaut riding a horse on mars"
+        with torch.autocast(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        ):  # for some reason device_type needs to be a string
+            # instead of an actual device
+            pipeline(prompt).images[0]
+
+
+def hf_main(args):
+    output_prefix = args.output_prefix
+    print("MODEL PATH:", args.input_directory)
+    print("OUTPUT PREFIX:", output_prefix)
+
+    from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+
+    model_config = AutoConfig.from_pretrained(args.input_directory)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.input_directory,
+        config=model_config,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+    )
+
+    logger.info("Serializing model")
+    logger.info("GPU: " + utils.get_gpu_name())
+    logger.info("PYTHON USED RAM: " + utils.get_mem_usage())
+
+    serialize_model(model, model_config, output_prefix)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.input_directory
+    ).save_pretrained(output_prefix)
+
+    if args.validate:
+        device = utils.get_device()
+        logger.info("Validating serialization")
+        model = load_model(
+            output_prefix,
+            AutoModelForCausalLM,
+            AutoConfig,
+            None,
+            device,
+            "float16",
+        ).eval()
+        # test generation
+        tokenizer = AutoTokenizer.from_pretrained(args.input_directory)
+        input_ids = tokenizer.encode(
+            "¡Hola! Encantado de conocerte. hoy voy a", return_tensors="pt"
+        ).to(device)
+        with torch.no_grad():
+            output = model.generate(
+                input_ids, max_new_tokens=50, do_sample=True
+            )
+        logger.info(
+            f"Test Output: {tokenizer.decode(output[0], skip_special_tokens=True)}"
+        )
+
+
+def main():
+    # usage: hf_serialization.py [-h] --model_type {transformers,diffusers} [--validate] input_directory output_prefix
+
+    parser = argparse.ArgumentParser(
+        description="An example script that uses Tensorizer to serialize an HF model to an output directory."
+    )
+    parser.add_argument(
+        "input_directory",
+        type=str,
+        help="Path to model directory or HF model ID",
+    )
+    parser.add_argument(
+        "output_prefix", type=str, help="Path to output directory"
+    )
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        choices=["transformers", "diffusers"],
+        required=True,
+        help="Framework used for the model",
+    )
+    parser.add_argument(
+        "--validate",
+        action="store_true",
+        help="Validate serialization by running a test inference",
+    )
+    args = parser.parse_args()
+
+    if args.model_type == "transformers":
+        hf_main(args)
+    elif args.model_type == "diffusers":
+        df_main(args)
+    else:
+        raise ValueError(
+            f"Unknown model type {args.model_type} (transformers or diffusers)"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/requirements.txt b/examples/requirements.txt
new file mode 100644
index 00000000..3f858ec4
--- /dev/null
+++ b/examples/requirements.txt
@@ -0,0 +1,3 @@
+transformers~=4.27.1
+diffusers==0.11.1
+accelerate==0.17.1
\ No newline at end of file
diff --git a/examples/serialize.py b/examples/serialize.py
new file mode 100644
index 00000000..ab59dbe0
--- /dev/null
+++ b/examples/serialize.py
@@ -0,0 +1,22 @@
+from transformers import AutoModelForCausalLM
+from tensorizer import TensorSerializer
+import torch
+
+model_ref = "EleutherAI/gpt-j-6B"
+# For less intensive requirements, swap above with the line below:
+# model_ref = "EleutherAI/gpt-neo-125m"
+model_name = model_ref.split("/")[-1]
+# Change this to your S3 bucket.
+s3_bucket = "bucket"
+s3_uri = f"s3://{s3_bucket}/{model_name}.tensors"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_ref,
+    revision="float16",
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+)
+
+serializer = TensorSerializer(s3_uri)
+serializer.write_module(model)
+serializer.close()
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..bb9e002b
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,41 @@
+[project]
+name = "tensorizer"
+version = "1.0.0"
+license = { text = "MIT License" }
+keywords = ["tensorizer", "machine learning", "serialization", "tensor", "pytorch"]
+authors = [
+  { name="CoreWeave" }
+]
+description = "A tool for fast PyTorch module, model, and tensor serialization + deserialization."
+readme = "README.md"
+requires-python = ">=3.8"
+dependencies = [
+    "torch>=1.9.0",
+    "numpy>=1.19.5",
+    "protobuf>=3.19.5",
+    "psutil>=5.9.4",
+    "boto3>=1.26.0"
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Development Status :: 5 - Production/Stable",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: System :: Distributed Computing",
+    "Topic :: Internet",
+    "Intended Audience :: Developers"
+]
+
+[project.urls]
+"Homepage" = "https://github.com/coreweave/tensorizer"
+
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+packages = ["tensorizer"]
+
+[tool.setuptools.package-data]
+tensorizer = ["tensors.proto"]
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index d8eabffc..36886d47 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 torch>=1.9.0
+numpy>=1.19.5
 protobuf>=3.19.5
-diffusers==0.11.1
-transformers==4.21.1
\ No newline at end of file
+psutil>=5.9.4
+boto3>=1.26.0
\ No newline at end of file
diff --git a/setup.py b/setup.py
deleted file mode 100644
index 76f9f0a8..00000000
--- a/setup.py
+++ /dev/null
@@ -1,28 +0,0 @@
-from setuptools import setup
-
-setup(
-    name='tensorizer',
-    version='0.0.1',
-    description='A tool for PyTorch Module, Model, and Tensor Serialization/Deserialization.',
-    url='https://github.com/coreweave/tensorizer',
-    author='CoreWeave',
-    license='MIT',
-    packages=['tensorizer'],
-    install_requires=[
-        'torch>=1.9.0',
-        'protobuf>=3.19.5',
-        'diffusers==0.11.1',
-        'transformers==4.21.1',
-    ],
-    classifiers=[
-        'Programming Language :: Python :: 3',
-        'License :: OSI Approved :: MIT License',
-        'Operating System :: OS Independent',
-    ],
-    python_requires='>=3.6',
-    entry_points={
-        'console_scripts': [
-            'tensorizer = tensorizer.tensorizer:main',
-        ],
-    },
-)
diff --git a/tensorizer/__init__.py b/tensorizer/__init__.py
index d1ca25d0..3060bce8 100644
--- a/tensorizer/__init__.py
+++ b/tensorizer/__init__.py
@@ -1,1 +1,10 @@
-from . import tensorizer
\ No newline at end of file
+from . import serialization, stream_io, utils
+from .serialization import *
+
+__all__ = [
+    *serialization.__all__,
+    "stream_io",
+    "utils",
+    "protobuf",
+    "tensors_pb2",
+]
diff --git a/tensorizer/_wide_pipes.py b/tensorizer/_wide_pipes.py
new file mode 100644
index 00000000..c5f2c476
--- /dev/null
+++ b/tensorizer/_wide_pipes.py
@@ -0,0 +1,161 @@
+import contextlib
+import functools
+import logging
+import threading
+import subprocess
+import sys
+
+# =============================================================================
+# From `pipe(7)` manpage:
+#
+# Pipe capacity
+# A pipe has a limited capacity. If the pipe is full, then a write(2) will
+# block or fail, depending on whether the O_NONBLOCK flag is set (see below).
+# Different implementations have different limits for the pipe capacity.
+#
+# Applications should not rely on a particular capacity: an application should
+# be designed so that a reading process consumes data as soon as it is
+# available, so that a writing process does not remain blocked.
+#
+# In Linux versions before 2.6.11, the capacity of a pipe was the same as the
+# system page size (e.g., 4096 bytes on i386). Since Linux 2.6.11, the pipe
+# capacity is 16 pages (i.e., 65,536 bytes in a system with a page size of
+# 4096 bytes). Since Linux 2.6.35, the default pipe capacity is 16 pages, but
+# the capacity can be queried and set using the fcntl(2) F_GETPIPE_SZ and
+# F_SETPIPE_SZ operations. See fcntl(2) for more information.
+#
+# =============================================================================
+# From `fcntl(2)` manpage:
+#
+# Changing the capacity of a pipe
+#
+# F_SETPIPE_SZ (int; since Linux 2.6.35)
+# Change the capacity of the pipe referred to by fd to be at least arg bytes.
+# An unprivileged process can adjust the pipe capacity to any value between the
+# system page size and the limit defined in /proc/sys/fs/pipe−max−size
+# (see proc(5)). Attempts to set the pipe capacity below the page size are
+# silently rounded up to the page size. Attempts by an unprivileged process to
+# set the pipe capacity above the limit in /proc/sys/fs/pipe−max−size yield the
+# error EPERM; a privileged process (CAP_SYS_RESOURCE) can override the limit.
+#
+# When allocating the buffer for the pipe, the kernel may use a capacity larger
+# than arg, if that is convenient for the implementation. (In the current
+# implementation, the allocation is the next higher power-of-two page-size
+# multiple of the requested size.) The actual capacity (in bytes) that is set
+# is returned as the function result.
+#
+# Attempting to set the pipe capacity smaller than the amount of buffer space
+# currently used to store data produces the error EBUSY.
+#
+# Note that because of the way the pages of the pipe buffer are employed when
+# data is written to the pipe, the number of bytes that can be written may be
+# less than the nominal size, depending on the size of the writes.
+#
+# F_GETPIPE_SZ (void; since Linux 2.6.35)
+# Return (as the function result) the capacity of the pipe referred to by fd.
+#
+# =============================================================================
+# Constant for `F_SETPIPE_SZ`, as Python's `fcntl` module doesn't have this
+# defined until Python 3.10.
+F_SETPIPE_SZ = 1031
+
+_logger = logging.getLogger(__name__)
+
+__all__ = ["get_max_pipe_size", "widen_pipe", "widen_new_pipes"]
+
+# No-op default implementations
+widen_new_pipes = contextlib.ExitStack
+
+
+def widen_pipe(_fileno):
+    pass
+
+
+@functools.lru_cache(maxsize=None)
+def get_max_pipe_size():
+    pipe_buf_sz = 1024 * 1024
+    if sys.platform != "win32":
+        # Read our max-fd-size, fall back to 1mb if invalid.
+        try:
+            with open("/proc/sys/fs/pipe-max-size", "r") as pipe_file:
+                pipe_buf_sz = int(pipe_file.read())
+        except IOError as e:
+            _logger.warning(
+                f"Could not read /proc/sys/fs/pipe-max-size: {e.strerror}"
+            )
+    else:
+        # Windows has no maximum pipe size,
+        # so 256 MiB is chosen completely arbitrarily.
+        pipe_buf_sz = 256 * 1024 * 1024
+    _logger.debug(f"pipe-max-size: {pipe_buf_sz}")
+    return pipe_buf_sz
+
+
+if sys.platform != "win32":
+    # Linux uses fcntl to resize an existing pipe.
+    import fcntl
+
+    def widen_pipe(fileno):
+        pipe_buf_sz = get_max_pipe_size()
+        try:
+            fcntl.fcntl(fileno, F_SETPIPE_SZ, pipe_buf_sz)
+        except PermissionError as e:
+            _logger.warning(
+                f"Couldn't fcntl F_SETPIPE_SZ to {pipe_buf_sz}: {e.strerror}"
+            )
+
+else:
+    # Windows cannot change the size of a pipe after creation,
+    # but it can set one's size during creation, so a context manager
+    # is used to temporarily modify the creation of all pipes.
+    _winapi = getattr(subprocess, "_winapi", None)
+    if _winapi is not None and hasattr(_winapi, "CreatePipe"):
+
+        class _LocalPipeSize(threading.local):
+            pipe_size = 0
+
+        _local = _LocalPipeSize()
+        _original_create_pipe = _winapi.CreatePipe
+        _pipe_routine_swap_mutex = threading.Lock()
+        _pipe_widening_threads = 0
+
+        def _create_wide_pipe(pipe_attrs, size):
+            # The subprocess module creates new anonymous pipes on Windows as:
+            # _winapi.CreatePipe(None, 0)
+            # Where the first argument is ignored,
+            # and the second is the pipe size (0 = default, usually 1 page).
+            # To change this without reimplementing all of subprocess.Popen,
+            # _winapi.CreatePipe itself is wrapped to override a size of 0
+            # with a default of our choosing.
+            #
+            # This function is thread-safe in the sense that other threads
+            # creating pipes while this is active will end up with
+            # unchanged results due to the override value being thread-local.
+            return _original_create_pipe(
+                pipe_attrs, _local.pipe_size if size == 0 else size
+            )
+
+        @contextlib.contextmanager
+        def widen_new_pipes():
+            global _pipe_widening_threads
+            # Thread-safe but not re-entrant.
+            # Thread safety in this function only matters if multiple threads
+            # try to separately invoke this context manager at the same time,
+            # which would only happen if multiple CURLStreamFiles were being
+            # opened in the same process at the same time. It is less important
+            # than _create_wide_pipe being thread-safe.
+            _local.pipe_size = get_max_pipe_size()
+            with _pipe_routine_swap_mutex:
+                _winapi.CreatePipe = _create_wide_pipe
+                _pipe_widening_threads += 1
+            try:
+                yield
+            finally:
+                with _pipe_routine_swap_mutex:
+                    _pipe_widening_threads -= 1
+                    if _pipe_widening_threads == 0:
+                        _winapi.CreatePipe = _original_create_pipe
+                del _local.pipe_size
+
+    else:
+        _logger.warning("Couldn't increase default pipe size.")
diff --git a/protobuf.py b/tensorizer/protobuf.py
similarity index 78%
rename from protobuf.py
rename to tensorizer/protobuf.py
index b5c97c6d..548dc988 100644
--- a/protobuf.py
+++ b/tensorizer/protobuf.py
@@ -3,7 +3,8 @@
 import numpy as np
 import tensors.tensors_pb2 as tensors_pb
 from tensors.tensors_pb2 import Tensor as TensorPb
-from typing import OrderedDict, Tuple, Union, BinaryIO
+from typing import OrderedDict as OrderedDictType, Tuple, Union, BinaryIO
+from collections import OrderedDict
 
 DtypePbs = {
     torch.float32: tensors_pb.DT_FLOAT32,
@@ -66,15 +67,18 @@
 }
 
 
-def serialize_tensor(t: Tensor, attribute: tensors_pb.AttributeType = None) -> \
-        tensors_pb.Tensor:
+def serialize_tensor(
+    t: Tensor, attribute: tensors_pb.AttributeType = None
+) -> tensors_pb.Tensor:
     assert isinstance(t, Tensor)
-    assert attribute is None or attribute in [tensors_pb.AT_PARAMETER,
-                                              tensors_pb.AT_BUFFER]
+    assert attribute is None or attribute in [
+        tensors_pb.AT_PARAMETER,
+        tensors_pb.AT_BUFFER,
+    ]
 
     extra_opts = {}
     if attribute is not None:
-        extra_opts = {'attr_type': attribute}
+        extra_opts = {"attr_type": attribute}
 
     return tensors_pb.Tensor(
         dtype=DtypePbs[t.dtype],
@@ -84,14 +88,15 @@ def serialize_tensor(t: Tensor, attribute: tensors_pb.AttributeType = None) -> \
     )
 
 
-def deserialize_tensor(t: tensors_pb.Tensor) -> \
-        Union[Tensor, Tuple[Tensor, 'tensors_pb.AttributeType']]:
+def deserialize_tensor(
+    t: tensors_pb.Tensor,
+) -> Union[Tensor, Tuple[Tensor, "tensors_pb.AttributeType"]]:
     mv = bytearray(t.data)
-    tensor = torch.as_tensor(np.ndarray.__new__(np.memmap,
-                                                t.shape,
-                                                dtype=PbNpyDtypes[t.dtype],
-                                                buffer=mv,
-                                                offset=0))
+    tensor = torch.as_tensor(
+        np.ndarray.__new__(
+            np.memmap, t.shape, dtype=PbNpyDtypes[t.dtype], buffer=mv, offset=0
+        )
+    )
     if t.HasField("attr_type"):
         return tensor, t.attr_type
     else:
@@ -106,22 +111,17 @@ def serialize_model(model: torch.nn.Module, file_stream: BinaryIO) -> None:
         for name, param in module.named_parameters(recurse=False):
             v = param.cpu().detach()
             param_attr = tensors_pb.Attribute(
-                name=name,
-                tensor=serialize_tensor(v, tensors_pb.AT_PARAMETER)
+                name=name, tensor=serialize_tensor(v, tensors_pb.AT_PARAMETER)
             )
             attributes.append(param_attr)
         for name, buffer in module.named_buffers(recurse=False):
             v = buffer.cpu().detach()
             buffer_attr = tensors_pb.Attribute(
-                name=name,
-                tensor=serialize_tensor(v, tensors_pb.AT_BUFFER)
+                name=name, tensor=serialize_tensor(v, tensors_pb.AT_BUFFER)
             )
             attributes.append(buffer_attr)
         module_attr = tensors_pb.Attribute(
-            name=module_name,
-            module=tensors_pb.Module(
-                attributes=attributes
-            )
+            name=module_name, module=tensors_pb.Module(attributes=attributes)
         )
         modules.append(module_attr)
     model_proto = tensors_pb.Module(  # models are just modules as attributes
@@ -135,7 +135,7 @@ def deserialize_model(model: torch.nn.Module, file_stream: BinaryIO) -> None:
     model_proto = tensors_pb.Module()
     model_proto.ParseFromString(file_stream.read())
 
-    modules: OrderedDict[str, torch.nn.Module] = OrderedDict()
+    modules: OrderedDictType[str, torch.nn.Module] = OrderedDict()
     for name, module in model.named_modules():
         modules[name] = module
 
@@ -144,8 +144,12 @@ def deserialize_model(model: torch.nn.Module, file_stream: BinaryIO) -> None:
         for attr in module_attr.module.attributes:
             if attr.tensor.HasField("attr_type"):
                 if attr.tensor.attr_type == tensors_pb.AT_PARAMETER:
-                    module._parameters[attr.name] = deserialize_tensor(attr.tensor)[0]
+                    module._parameters[attr.name] = deserialize_tensor(
+                        attr.tensor
+                    )[0]
                 elif attr.tensor.attr_type == tensors_pb.AT_BUFFER:
-                    module._buffers[attr.name] = deserialize_tensor(attr.tensor)[0]
+                    module._buffers[attr.name] = deserialize_tensor(
+                        attr.tensor
+                    )[0]
                 else:
                     raise ValueError("Unknown attribute type")
diff --git a/tensorizer/serialization.py b/tensorizer/serialization.py
new file mode 100644
index 00000000..c10f6558
--- /dev/null
+++ b/tensorizer/serialization.py
@@ -0,0 +1,1122 @@
+##############################################################################
+# serialization.py                                                   Wes Brown
+# Fast torch module/model serialization/deserialization     (c) 2023 Coreweave
+##############################################################################
+
+# try to import UNIX only dependencies
+try:
+    import fcntl
+    import resource
+except ImportError:
+    fcntl = None
+    resource = None
+
+from enum import Enum
+import collections.abc
+import io
+import os
+import tensorizer.stream_io as stream_io
+import tensorizer.utils as utils
+import mmap
+import numpy
+import struct
+import time
+import torch
+import typing
+import logging
+import tempfile
+import hashlib
+import zlib
+import ctypes
+
+if torch.cuda.is_available():
+    cudart = torch.cuda.cudart()
+else:
+    cudart = None
+
+from collections import OrderedDict
+from typing import Optional, Tuple, Union, List, Iterator, Dict, Callable, Any
+
+lz4 = None
+
+__all__ = ["TensorSerializer", "TensorDeserializer", "TensorType"]
+
+# Setup logger
+logger = logging.getLogger(__name__)
+
+
+# Whether the tensor is a parameter or a buffer on the model.
+class TensorType(Enum):
+    PARAM = 0
+    BUFFER = 1
+    STATE_DICT = 2
+
+
+TENSORIZER_VERSION = 1
+TENSORIZER_MAGIC = b"|TZR|"
+
+
+class TensorEntry(typing.TypedDict):
+    name: str
+    type: TensorType
+    offset: int
+    data_offset: int
+    data_length: int
+    dtype: str
+    shape: List[int]
+
+
+class HashType(Enum):
+    CRC32 = 0
+    SHA256 = 1
+
+
+class TensorHash(typing.TypedDict):
+    type: HashType
+    hash: bytes
+
+
+def _convert_dtype_to_numpy(dtype: Union[numpy.dtype, str, torch.dtype]):
+    if isinstance(dtype, numpy.dtype):
+        return dtype
+    elif isinstance(dtype, str):
+        return numpy.dtype(dtype)
+    elif isinstance(dtype, torch.dtype):
+        # Converted from PyTorch's own mapping used in their testing:
+        # https://github.com/pytorch/pytorch/blob/v2.0.0/torch/testing/_internal/common_utils.py#L1009
+        numpy_dtype = {
+            torch.bool: numpy.bool_,
+            torch.uint8: numpy.uint8,
+            torch.int8: numpy.int8,
+            torch.int16: numpy.int16,
+            torch.int32: numpy.int32,
+            torch.int64: numpy.int64,
+            torch.float16: numpy.float16,
+            torch.float32: numpy.float32,
+            torch.float64: numpy.float64,
+            torch.complex64: numpy.complex64,
+            torch.complex128: numpy.complex128,
+        }.get(dtype)
+
+        if numpy_dtype is None:
+            raise TypeError(
+                "The provided torch.dtype class could not"
+                " be converted to a numpy.dtype"
+            )
+        else:
+            # Convert it from a dtype class to a dtype object instance
+            return numpy.dtype(numpy_dtype)
+    else:
+        raise TypeError("Could not interpret provided dtype as a numpy.dtype")
+
+
+class TensorDeserializer(collections.abc.Mapping):
+    """
+    Given a file-like object for read, deserialize tensors to a state_dict or
+    a torch.nn.Module.
+
+    Args:
+        file_obj: A file-like object to read from. It can also be a string
+            representing a path to a file or an HTTP/HTTPS/S3 URI.
+        device: The device to load the tensors to.
+        filter_func: A function (tensor_name: str) -> bool that returns True
+            if a tensor should be loaded, or False if it should be skipped.
+            If None, all tensors are loaded.
+        dtype: The dtype to load the tensors as. If None, the dtype will be
+            inferred from the file.
+        lazy_load: If True, tensors will be loaded and cached when keys are
+            accessed. If False, all tensors will be loaded into memory up
+            front.
+        plaid_mode: If True, tensors will be loaded extremely fast into the
+            target device. This is only supported on CUDA devices, and the
+            buffers are going to be inconsistent due to the extreme
+            naughtiness of reusing a backing buffer. This is only recommended
+            for use with inference, and not training.
+    """
+
+    def __init__(
+        self,
+        file_obj: Union[io.BufferedIOBase, io.RawIOBase, typing.BinaryIO, str],
+        device: Union[torch.device, str, None] = None,
+        filter_func: Optional[Callable[[str], Union[bool, Any]]] = None,
+        dtype: Union[numpy.dtype, str, None] = None,
+        *,
+        lazy_load: bool = False,
+        plaid_mode: bool = False,
+    ):
+        if isinstance(file_obj, str):
+            self._file = stream_io.open_stream(file_obj, "rb")
+        else:
+            self._mode_check(file_obj)
+            self._file = file_obj
+        self.total_compressed_tensor_bytes = 0
+        self.read_bytes = 0
+
+        # If device is None, use the current device, otherwise use the given
+        # device.
+        device = utils.get_device() if device is None else torch.device(device)
+        self._device = device
+
+        # If dtype is not None, convert all tensors to this dtype if possible.
+        if dtype is None:
+            self._dtype = None
+        else:
+            self._dtype = _convert_dtype_to_numpy(dtype)
+
+        self._metadata: Dict[str, TensorEntry] = {}
+
+        if plaid_mode and (
+            not torch.cuda.is_available() or self._device.type == "cpu"
+        ):
+            raise ValueError("Plaid mode requires CUDA")
+
+        # plaid_mode implies lazy_load
+        self._lazy_load = lazy_load or plaid_mode
+
+        # Read the magic
+        magic = self._file.read(5)
+        if magic != TENSORIZER_MAGIC:
+            raise ValueError("Not a tensorizer file")
+
+        # Read the version
+        version = struct.unpack("<I", self._file.read(4))[0]
+
+        # Check the version
+        if version > TENSORIZER_VERSION:
+            raise ValueError(
+                f"This {TENSORIZER_VERSION} version cannot read"
+                f"file versioned {version}."
+            )
+
+        # Skip 32-byte hash (unused)
+        self._file.read(32)
+
+        # Read the total size of the file
+        self.total_file_bytes = struct.unpack("<Q", self._file.read(8))[0]
+
+        # Read total size of tensor data
+        self.total_tensor_bytes = struct.unpack("<Q", self._file.read(8))[0]
+
+        # Read the number of tensors
+        self._tensors = struct.unpack("<Q", self._file.read(8))[0]
+
+        # Read the metadata index of tensors. This is a list of offsets into the
+        # file where the per-tensor data is stored. filter_func is a test that
+        # determines the tensor names to read. If filter_func is None,
+        # all tensors are read.
+        self._load_metadatas(filter_func)
+
+        self._plaid_mode: bool = plaid_mode
+        self._prior_key: Optional[str] = None
+
+        # We calculate the total tensor bytes here so that we can use mmap,
+        # based on the total size of the tensors that we're going to read,
+        # filtered by the filter_func.
+        self.total_tensor_bytes = 0
+        self._largest_tensor_bytes = 0
+        for name, metadata in self._metadata.items():
+            self.total_tensor_bytes += metadata["data_length"]
+            if metadata["data_length"] > self._largest_tensor_bytes:
+                self._largest_tensor_bytes = metadata["data_length"]
+
+        # Allocate the buffer for the tensors. If we're not in lazy_load mode,
+        # we'll allocate a single buffer for all the tensors. Otherwise, we'll
+        # allocate a buffer for the largest tensor.
+        self._buffer_addr = None
+        self._is_memory_pinned = False
+        if not lazy_load and not plaid_mode:
+            start_allocate = time.time()
+
+            # Check if our platform supports mmap.MAP_ANONYMOUS and
+            # mmap.MAP_PRIVATE
+            mmap_flags = 0
+            mmap_flags |= getattr(mmap, "MAP_PRIVATE", 0)
+            mmap_flags |= getattr(mmap, "MAP_ANONYMOUS", 0)
+            mmap_args = {"flags": mmap_flags} if mmap_flags else {}
+            self._buffer = mmap.mmap(-1, self.total_tensor_bytes, **mmap_args)
+
+            # If we're on CUDA, we register the buffer with CUDA so that
+            # it is pinned. This allows for Torch to internally use
+            # cudaMemcpyAsync.
+            if self._device.type == "cuda":
+                # We need to use ctypes to get the address of the buffer
+                # because mmap.mmap doesn't expose the buffer address.
+                tb = ctypes.c_char * self.total_tensor_bytes
+                ctb = tb.from_buffer(self._buffer)
+                self._buffer_addr = ctypes.addressof(ctb)
+                del ctb  # don't leave an open exported pointer into the mmap
+
+                # Register the buffer with CUDA
+                cudart.cudaHostRegister(
+                    self._buffer_addr, self.total_tensor_bytes, 0
+                )
+                self._is_memory_pinned = True
+            end_allocate = time.time()
+            tensor_bytes_str = utils.convert_bytes(self.total_tensor_bytes)
+            logger.info(
+                f"Allocated {tensor_bytes_str} "
+                f"for {len(self._metadata)} tensors "
+                f"in {end_allocate - start_allocate:0.4f}"
+            )
+        else:
+            self._buffer = bytearray(self._largest_tensor_bytes)
+
+        # The number of bytes we've allocated so far. Tensors may be read
+        # from the file in any order, so we need to keep track of how much
+        # we've used so far so that we can index into the buffer correctly.
+        self._allocated = 0
+
+        # Our cache of tensors. This is a dict of name -> tensor. If lazy_load
+        # is True, then the tensors are not loaded until they are accessed.
+        self._cache: typing.OrderedDict[str, Union[torch.Tensor, None, bool]]
+
+        # The offset in the file where the tensor data begins.
+        self._tensors_begin = self._file.tell()
+
+        if not lazy_load:
+            # If we're not in lazy_load mode, we populate the cache with all
+            # the tensors.
+            self._cache = self._generate_state_dict()
+        else:
+            # We populate the cache with None values so that we can
+            # differentiate between tensors that have not been loaded yet
+            # and tensors that are not present in the file.
+            self._cache = OrderedDict.fromkeys(self._metadata.keys())
+
+    @staticmethod
+    def _mode_check(file_obj: io.IOBase) -> None:
+        try:
+            readable = file_obj.readable()
+        except AttributeError:
+            # If file_obj doesn't implement the full io.IOBase interface
+            # and checking is not possible, assume it's fine
+            readable = True
+        if isinstance(file_obj, io.TextIOBase) or not readable:
+            mode = getattr(file_obj, "mode", "")
+            raise ValueError(
+                "TensorSerializer's file_obj must be readable "
+                'and in binary mode (mode="rb"{})'.format(
+                    mode and f', current mode="{mode}"'
+                )
+            )
+
+    def __del__(self):
+        self.close()
+
+    def close(self):
+        # Don't throw an attribute error if these aren't defined yet,
+        # e.g. if __init__ threw an error before defining both
+        buffer = getattr(self, "_buffer", None)
+        if buffer is not None:
+            if self._is_memory_pinned and self._buffer_addr:
+                cudart.cudaHostUnregister(self._buffer_addr)
+                self._is_memory_pinned = False
+            if self._device.type != "cpu" and hasattr(buffer, "close"):
+                # Don't close the mmap buffer for CPU tensors because
+                # it would free their memory out from under them.
+                # For GPU tensors we just need to wait until
+                # they are transferred out of RAM.
+                buffer.close()
+        if getattr(self, "_file", None) is not None:
+            self._file.close()
+
+    def _read_string(self, io_obj=None):
+        """
+        Read a string from the file.
+        """
+        if io_obj is None:
+            io_obj = self._file
+
+        length = struct.unpack("<H", io_obj.read(2))[0]
+        return io_obj.read(length).decode("utf-8")
+
+    def _read_dtype(self, io_obj=None):
+        """
+        Read a dtype from the file.
+        """
+        if io_obj is None:
+            io_obj = self._file
+
+        length = struct.unpack("<B", io_obj.read(1))[0]
+        return io_obj.read(length).decode("utf-8")
+
+    def _read_metadata(self, metadata_stream: io.BytesIO) -> TensorEntry:
+        name = self._read_string(metadata_stream)
+        tensor_type = TensorType(
+            struct.unpack("<B", metadata_stream.read(1))[0]
+        )
+        dtype = self._read_dtype(metadata_stream)
+        shape_len = struct.unpack("<B", metadata_stream.read(1))[0]
+        shape = self._read_shapes(
+            metadata_stream.read(shape_len * 4), shape_len
+        )
+        offset = struct.unpack("<Q", metadata_stream.read(8))[0]
+        data_offset = struct.unpack("<Q", metadata_stream.read(8))[0]
+        data_length = struct.unpack("<Q", metadata_stream.read(8))[0]
+        return TensorEntry(
+            name=name,
+            type=tensor_type,
+            offset=offset,
+            data_offset=data_offset,
+            data_length=data_length,
+            dtype=dtype,
+            shape=shape,
+        )
+
+    def _load_metadatas(
+        self, filter_func: Optional[Callable[[str], Union[bool, Any]]]
+    ):
+        """
+        Read the metadata of tensors into self._metadata.
+        """
+
+        # Read metadata size.
+        self._metadata_size = struct.unpack("<Q", self._file.read(8))[0]
+        metadata_encoded = self._file.read(self._metadata_size)
+        # Turn the metadata into a stream.
+        metadata_stream = io.BytesIO(metadata_encoded)
+
+        for i in range(self._tensors):
+            metadata = self._read_metadata(metadata_stream)
+            if filter_func is None or filter_func(metadata["name"]):
+                self._metadata[metadata["name"]] = metadata
+
+    @staticmethod
+    def _read_shapes(obj, num_elems) -> List[int]:
+        """
+        Read the tensor shapes.
+        """
+        return list(struct.unpack(f"<{num_elems}I", obj))
+
+    @property
+    def total_bytes_read(self) -> int:
+        if self._file.closed:
+            return self.total_tensor_bytes
+        else:
+            return self._file.tell()
+
+    def __getitem__(self, name) -> torch.nn.Parameter:
+        if self._plaid_mode:
+            # In plaid_mode, we only have one valid tensor at a time, so
+            # we need to make sure that prior tensors are not accessed.
+            if self._prior_key is not None and self._prior_key != name:
+                self._cache[self._prior_key] = False
+            if self._cache[name] is False:
+                raise RuntimeError(
+                    f"Tensor {name} already overwritten in " "plaid_mode"
+                )
+            self._prior_key = name
+
+        if name in self._cache and self._cache[name] is not None:
+            return self._cache[name]
+
+        # If we're in lazy_load mode, we populate the cache with the
+        # tensor data and then convert it to a torch parameter. Most
+        # of the time, access patterns are front to back, so seeking
+        # forward in a stream works well even for HTTP/HTTPS streams.
+        if name in self._metadata:
+            self._file.seek(self._metadata[name]["offset"])
+            tensor_arr = next(self.read_tensors(num_tensors=1))[3]
+            self._cache[name] = self._to_torch_parameter(tensor_arr)
+            return self._cache[name]
+        else:
+            raise KeyError(f"Tensor {name} not found")
+
+    # To implement collections.abc.Mapping, this class needs to define:
+    # 1. __getitem__(key)
+    # 2. __iter__()
+    # 3. __len__()
+    #
+    # It then inherits efficient default implementations for:
+    # 1. get(key)
+    # 2. keys()
+    # 3. items()
+    # 4. values()
+    #
+    # It also inherits __contains__(key), but it isn't efficient.
+    # The default __contains__ implementation uses __getitem__,
+    # so it loads tensor data unnecessarily. Instead, we can
+    # check self._metadata with a simple dict lookup.
+    #
+    # Significant features of the defaults:
+    # - values() -> collections.abc.ValuesView:
+    #   - Doesn't need to load all the data up front
+    #   - Returns a collection with a lazy __iter__
+    #     - Only defers to our __getitem__ exactly when required
+    #   - Implements __contains__ by iterating over itself with linear search
+    # - items() -> collections.abc.ItemsView:
+    #   - Also doesn't need to load data up front
+    #   - Implements __iter__ as basically zip(keys(), values())
+    #   - Implements __contains__ with __getitem__ since it knows where to
+    #     look and needs the data anyway
+    #
+    # In summary, ignoring the (still efficient) __contains__, these are
+    # simply lazy iterables over the respective elements from this parent
+    # mapping.
+
+    def __iter__(self):
+        # iter() on a mapping returns an iterator of only keys
+        yield from self._metadata
+
+    def __len__(self):
+        return len(self._metadata)
+
+    def __contains__(self, key: str):
+        return key in self._metadata
+
+    def keys(self):
+        # We override keys() because dict_keys can be slightly more efficient
+        # than an extra collections.abc.KeysView wrapper.
+        #
+        # Technically this makes mapping.keys().mapping invalid on
+        # Python 3.10+ but it is not intended to be supported anyway, so treat
+        # it as not implemented.
+        return self._metadata.keys()
+
+    @staticmethod
+    def _decode_hashes(b: bytes) -> List[TensorHash]:
+        """
+        Decode the hashes from given bytes.
+        """
+        hashes: List[TensorHash] = []
+
+        # Read the number of hashes.
+        num_hashes = struct.unpack("<B", b[0:1])[0]
+
+        hash_idx = 1
+        # Read the hashes.
+        for i in range(num_hashes):
+            # Read the hash type.
+            hash_type = struct.unpack("<B", b[hash_idx : hash_idx + 1])[0]
+            # Read the size of the hash.
+            hash_size = struct.unpack("<B", b[hash_idx + 1 : hash_idx + 2])[0]
+            # Read the hash.
+            hash_begin = hash_idx + 2
+            hash_end = hash_begin + hash_size
+            hash_bytes: bytes = b[hash_begin:hash_end]
+            # Add the hash to the list.
+            hash_entry = TensorHash(
+                type=HashType(hash_type),
+                hash=hash_bytes,
+            )
+            hash_idx = hash_end
+            hashes.append(hash_entry)
+
+        return hashes
+
+    def read_tensors(
+        self,
+        filter_func: Optional[Callable[[str], Union[bool, Any]]] = None,
+        num_tensors: int = -1,
+    ) -> Iterator[Tuple[int, int, str, numpy.ndarray]]:
+        """
+        A generator that deserializes tensors and returns the `module_idx`,
+        `tensor_type`, parameter/buffer `name`, and the numpy `arr` that
+        represents the tensor.
+
+        Note that this function does not seek to the beginning of the tensor
+        data. It assumes that the file pointer is already at the beginning
+        of the tensor data that it should read.
+
+        It will read `num_tensors` tensors from the file, or all tensors
+        if `num_tensors` is -1.
+
+        The generator yields tuples of the form:
+            (module_idx, tensor_type, name, arr)
+
+        Args:
+            filter_func: A function that takes a tensor name and returns
+                True if the tensor should be returned, False otherwise.
+            num_tensors: The number of tensors to read. If -1, all tensors
+                will be read. If the zero-byte header is encountered before
+                `num_tensors` tensors are read, the generator will stop
+                yielding values.
+        """
+
+        try:
+            tensors_read = 0
+            while num_tensors == -1 or tensors_read < num_tensors:
+                header_sz = struct.unpack("<Q", self._file.read(8))[0]
+                if header_sz == 0:
+                    break
+                # We read the entire header into memory rather than reading
+                # it piecewise to avoid the overhead of many small reads,
+                # especially for network streams.
+                headers = self._file.read(header_sz - 8)
+                header_len = len(headers)
+
+                module_idx = struct.unpack("<H", headers[0:2])[0]
+
+                tensor_type = TensorType(struct.unpack("<B", headers[2:3])[0])
+
+                # Read the name.
+                name_sz = struct.unpack("<H", headers[3:5])[0]
+                idx = name_sz + 5
+                name_bytes = headers[5:idx]
+                name: str = name_bytes.decode("utf-8")
+
+                # Read the dtype of the tensor.
+                dtype_len = struct.unpack("<B", headers[idx : idx + 1])[0]
+                dtype_end = idx + dtype_len + 1
+                dtype = headers[idx + 1 : dtype_end]
+
+                # Read the shape amount, according to the serialized format.
+                # The shape length is 1 byte after the dtype end.
+                shape_len = struct.unpack(
+                    "<B", headers[dtype_end : dtype_end + 1]
+                )[0]
+                # The shape elements are <I, so we read 4 bytes. _read_shapes
+                # takes in the header object and the number of elements in
+                # the shape.
+                #
+                # The amount of bytes for the shape is 4 * number of elements
+                # in the shape. so, we need to read 4 * shape_len bytes after
+                # the dtype end + 1 byte for the shape length. sort of
+                # convoluted, but it works.
+                shape_begin = dtype_end + 1
+                shape_end = shape_begin + (4 * shape_len)
+                shape_list = self._read_shapes(
+                    headers[shape_begin:shape_end],
+                    shape_len,
+                )
+
+                # Read our hashes in. We need to read the hashes size, then
+                # read the hash bytes.
+                #
+                # TODO: Actually verify the hashes on request.
+                hashes_begin = shape_end
+                hashes_sz_slice = headers[hashes_begin : hashes_begin + 2]
+                hashes_sz = struct.unpack("<H", hashes_sz_slice)[0]
+                hashes_end = hashes_begin + hashes_sz
+
+                hashes_slice = headers[hashes_begin + 2 : hashes_end]
+                hashes = self._decode_hashes(hashes_slice)
+
+                # Finally, get the tensor data length.
+                data_length = struct.unpack("<q", headers[header_len - 8 :])[0]
+
+                # Check if the name is in our pre-filtered list of keys
+                # from the class-level filter_func, and then verify
+                # that it passes the method-level filter_func.
+                # Skip it if it fails either check.
+                if name not in self.keys() or (
+                    filter_func is not None and not filter_func(name)
+                ):
+                    self._file.seek(data_length, io.SEEK_CUR)
+                    continue
+
+                # We use memoryview to avoid copying the data.
+                mv: memoryview
+                if not self._plaid_mode and not self._lazy_load:
+                    # In default mode, we've already allocated all the
+                    # memory we need in a single buffer that contains
+                    # all the tensors. We just need to slice out the
+                    # memoryview for the current tensor.
+                    mv = memoryview(self._buffer)[
+                        self._allocated : self._allocated + data_length
+                    ]
+                    self._file.readinto(mv)
+                    self._allocated += data_length
+                elif self._plaid_mode:
+                    # In plaid_mode, we don't allocate a buffer, we just
+                    # reuse the same one. This is a filthy hack, as we're
+                    # overwriting the buffer contents that is used to back
+                    # the prior tensor. This works because we don't use
+                    # the buffer contents after we yield the tensor, which
+                    # is loaded straight into the GPU memory.
+                    mv = memoryview(self._buffer)
+                    self._file.readinto(mv[:data_length])
+                else:
+                    # In lazy_load mode, we allocate a new buffer for each
+                    # tensor. This is a bit slower, but it's the only way
+                    # to support lazy loading.
+                    buffer = bytearray(data_length)
+                    mv = memoryview(buffer)
+                    self._file.readinto(mv)
+
+                # Convert the memoryview to a numpy array. We use the
+                # memmap class to avoid copying the data.
+                arr = numpy.ndarray.__new__(
+                    numpy.memmap,
+                    shape_list,
+                    dtype=dtype,
+                    buffer=mv,
+                    offset=0,
+                )
+                tensors_read += 1
+                yield module_idx, tensor_type, name, arr
+        except EOFError:
+            return
+
+    def _to_torch_parameter(
+        self, arr: Union[numpy.ndarray, torch.nn.Parameter]
+    ) -> torch.nn.Parameter:
+        """
+        Convert a numpy array to a torch.nn.Parameter on a device, forcing
+        gradient when appropriate. We also handle torch.nn.Parameter objects in
+        a passthrough manner.
+        """
+        if isinstance(arr, torch.nn.Parameter):
+            arr.data = arr.data.to(self._device)
+            if arr.grad is not None:
+                arr.grad = arr.grad.to(self._device)
+            return arr
+
+        if (
+            self._dtype is not None
+            and arr.dtype != "bool"
+            and arr.dtype != self._dtype
+        ):
+            arr = arr.astype(self._dtype)
+        gradient = arr.dtype.kind in ("f", "c")
+
+        return torch.nn.Parameter(
+            torch.from_numpy(arr).to(self._device), requires_grad=gradient
+        )
+
+    def _generate_state_dict(self) -> OrderedDict:
+        """
+        Load the tensors in this Tensorizer object into a state_dict. This
+        is used to populate the cache in non-lazy_load cases.
+        """
+        if self._file.closed:
+            raise IOError("IO closed, instantiate if you want to load again.")
+
+        d = OrderedDict()
+        for idx, typ, name, arr in self.read_tensors():
+            d[name] = self._to_torch_parameter(arr)
+        self.total_tensor_bytes = self._file.tell()
+        self._file.close()
+        return d
+
+    def load_into_module(
+        self,
+        m: torch.nn.Module,
+        filter_func: Optional[Callable[[str], Union[bool, Any]]] = None,
+    ) -> int:
+        """
+        Given `m`, a torch.nn.Module, load the associate tensors in this
+        Tensorizer object into the `torch.nn.Module`. Returns the number of
+        tensors loaded into the module.
+
+        Args:
+            m: The module to load the tensors into.
+            device: The device to load the tensors onto.
+            dtype: The dtype to load the tensors as. Defaults to None, which
+                means the dtype is not changed from the serialized dtype.
+            filter_func: A function (tensor_name: str) -> bool that returns
+                True if a tensor should be loaded, or False if it should be
+                skipped.
+        """
+        modules: typing.OrderedDict[str, torch.nn.Module] = OrderedDict()
+
+        for name, module in m.named_modules():
+            modules[name] = module
+
+        tensor_ct = 0
+
+        for name in self.keys():
+            if filter_func is not None and not filter_func(name):
+                continue
+            obj_path, attr = name.rsplit(".", 1)
+            module: torch.nn.Module = modules[obj_path]
+            entry = self._metadata[name]
+            param = self._to_torch_parameter(self.get(name))
+            if entry["type"] is TensorType.PARAM:
+                module.register_parameter(attr, param)
+            elif entry["type"] is TensorType.BUFFER:
+                module.register_buffer(attr, param)
+            elif entry["type"] is TensorType.STATE_DICT:
+                raise NotImplementedError(
+                    "This was serialized using the write_state_dict() method,"
+                    " and cannot be loaded using the load_tensors() method."
+                    " Use the TensorDeserializer object directly as a"
+                    " state_dict mapping instead."
+                )
+            tensor_ct += 1
+
+        self._file.close()
+        return tensor_ct
+
+
+class TensorSerializer:
+    """
+    Given a file-like object or path, serialize tensors from a torch.nn.Module
+    to it.
+
+    Args:
+        file_obj: A file-like object or path to a file to write to. The path
+            can be a S3 URI.
+        compress_tensors: If True, compress the tensors using lz4. This
+            exists as an internal curiousity as it doesn't seem to make
+            much of a difference in practice.
+    """
+
+    def __init__(
+        self,
+        file_obj: Union[
+            io.BufferedIOBase,
+            io.RawIOBase,
+            typing.BinaryIO,
+            str,
+            bytes,
+            os.PathLike,
+            int,
+        ],
+        compress_tensors: bool = False,
+    ) -> None:
+        if isinstance(file_obj, (str, bytes, os.PathLike, int)):
+            self._file = stream_io.open_stream(file_obj, "wb+")
+        else:
+            self._mode_check(file_obj)
+            self._file = file_obj
+        self._tensors = 0
+        self._idx = 0
+        self.total_tensor_bytes = 0
+        self.total_compressed_tensor_bytes = 0
+        self.compress_tensors = compress_tensors
+        if self.compress_tensors:
+            import lz4.frame
+
+            self.lz4_frame = lz4.frame
+        else:
+            self.lz4_frame = None
+
+        # Write our magic bytes.
+        self._file.write(TENSORIZER_MAGIC)
+
+        # Write the version number.
+        self._file.write(struct.pack("<I", TENSORIZER_VERSION))
+
+        # Reserve 32 bytes for the hash. (Unused for now)
+        self._hash_loc = self._file.tell()
+        self._file.write(struct.pack("<Q", 0) * 4)
+
+        # Reserve 8 bytes for the total size of the file.
+        self._size_loc = self._file.tell()
+        self._file.write(struct.pack("<Q", 0))
+
+        # Reserve 8 bytes for the total size of tensor data.
+        self._tensor_size_loc = self._file.tell()
+        self._file.write(struct.pack("<Q", 0))
+
+        # Reserve the next 8 bytes for the total number of tensors.
+        self._tensor_ct_loc = self._file.tell()
+        self._file.write(struct.pack("<Q", 0))
+
+        # Reserve 256kb for metadata.
+        metadata_size = 256 * 1024
+        self._file.write(struct.pack("<Q", metadata_size))
+        self._metadata_loc = self._file.tell()
+        self._file.write(struct.pack("<Q", 0) * (metadata_size // 8))
+        self._metadata_cur = self._metadata_loc
+        self._metadata_end = self._metadata_loc + metadata_size
+
+        self._tensor_index: List[TensorEntry] = []
+
+    @staticmethod
+    def _mode_check(file_obj: io.IOBase) -> None:
+        try:
+            read_write = file_obj.writable() and file_obj.readable()
+        except AttributeError:
+            # If file_obj doesn't implement the full io.IOBase interface
+            # and checking is not possible, assume it's fine
+            read_write = True
+        if isinstance(file_obj, io.TextIOBase) or not read_write:
+            mode = getattr(file_obj, "mode", "")
+            raise ValueError(
+                "TensorSerializer's file_obj must be readable, writable, "
+                'and in binary mode (mode="wb+"{})'.format(
+                    mode and f', current mode="{mode}"'
+                )
+            )
+
+    def __del__(self):
+        if getattr(self, "_file", None) is not None:
+            self._file.close()
+
+    @staticmethod
+    def _dump_shape(obj) -> bytes:
+        """
+        Returns shape of the tensor as an encoded byte string.
+        """
+        bstr = struct.pack("<B", len(obj))
+        for i in obj:
+            bstr += struct.pack("<I", i)
+        return bstr
+
+    def close(self) -> None:
+        """
+        Finalizes the serialization and closes the file.
+        """
+        self._sync_prologue_state()
+
+        final_sz = self._file.tell()
+        self._file.close()
+        logger.info(f"Tensors completed serializing to {final_sz} bytes")
+        if self.compress_tensors:
+            compression_ratio = (
+                self.total_tensor_bytes / self.total_compressed_tensor_bytes
+            )
+            logger.info(f"Uncomp'd bytes: {self.total_tensor_bytes}")
+            logger.info(f"Comp'd bytes: {self.total_compressed_tensor_bytes}")
+            logger.info(f"Ratio: {compression_ratio:.2f}")
+
+    def _sync_prologue_state(self):
+        """
+        This is called after the tensor has been written to the file, and
+        ensures that the file is in a consistent state.
+        :return:
+        """
+        curr = self._file.tell()
+
+        # Write our zero-length field, that indicates that this is the last
+        # tensor. This will be overwritten if another tensor is written.
+        self._file.write(struct.pack("<Q", 0))
+        # Write the total number of tensors.
+        self._file.seek(self._tensor_ct_loc)
+        self._file.write(struct.pack("<Q", self._tensors))
+        # Write our total file size.
+        self._file.seek(self._size_loc)
+        self._file.write(struct.pack("<Q", curr))
+        # Write the total bytes of tensor data written.
+        self._file.seek(self._size_loc)
+        self._file.write(struct.pack("<Q", self.total_tensor_bytes))
+
+        # Reset our file pointer to the end of the file, minus the zero-length
+        # field.
+        self._file.seek(curr)
+
+    def write_tensor(
+        self,
+        idx,
+        name,
+        tensor_type: TensorType,
+        tensor: Union[torch.Tensor, numpy.ndarray],
+    ) -> None:
+        """
+        Serializes a tensor, laying things out so that it can be read in three
+        calls from the input -- once for the size, once for the header, and
+        once for the tensor itself.
+
+        Args:
+            idx: The index of the tensor in the module.
+            name: The name of the tensor.
+            tensor_type: The type of the tensor. This is used to determine
+                how to interpret the tensor.
+            tensor: The tensor to serialize.
+
+        Serialization format:
+
+         { uint64                           header_sz,
+           uint16                           module_idx,
+           uint8                            type,
+           uint16                           name_sz,
+           []char                           name,
+           uint8                            dtype_sz,
+           []char                           dtype_str,
+           []{uint8                         shape_elem_type,
+              union{uint8, uint16, uint32}  shape_elem_sz,
+             }                              shape_elements,
+           uint16                           hashes_sz,
+           uint8                            num_hashes,
+           []{uint8                         hash_type,
+              uint8                         hash_sz,
+              []char                        hash_str,
+             }                              hashes,
+           uint64                           tensor_sz,
+           []byte                           tensor }
+        """
+
+        if isinstance(tensor, torch.Tensor):
+            tensor = tensor.cpu().detach().numpy()
+
+        if len(str(tensor.dtype)) >= 256:
+            raise ValueError("dtype length should be less than 256")
+
+        # Reserve room for our tensor header size.
+        ds_header_begin = self._file.tell()
+        self._file.write(struct.pack("<Q", 0))
+
+        # Module index.
+        self._file.write(struct.pack("<H", idx))
+
+        # Whether this is a parameter or a buffer
+        self._file.write(struct.pack("<B", tensor_type.value))
+
+        # Parameter/buffer name
+        name_bytes = bytes(name, "utf-8")
+        self._file.write(struct.pack("<H", len(name_bytes)))
+        self._file.write(name_bytes)
+
+        # Write out our tensor dtype
+        dtype_bytes = bytes(tensor.dtype.str, "utf-8")
+        dtype_len = len(dtype_bytes)
+        self._file.write(struct.pack("<B", dtype_len))
+        self._file.write(dtype_bytes)
+
+        # ... and shape
+        shape_bytes = self._dump_shape(tensor.shape)
+        self._file.write(shape_bytes)
+
+        # Reserve room for our tensor hashes.
+        hash_loc = self._file.tell()
+        # Reserve the length of the hash data structures.
+        self._file.write(struct.pack("<H", 0))
+        # Write the number of hashes we're going to write.
+        self._file.write(struct.pack("<B", 2))
+        # Reserve room for the hashes.
+        # .. first, CRC32.
+        self._file.write(struct.pack("<B", HashType.CRC32.value))
+        # ... length of CRC32.
+        self._file.write(struct.pack("<B", 4))
+        # ... and reserve for the actual CRC32.
+        crc32_loc = self._file.tell()
+        self._file.write(struct.pack("<I", 0))
+        # .. second, SHA256.
+        self._file.write(struct.pack("<B", HashType.SHA256.value))
+        # ... length of SHA256.
+        self._file.write(struct.pack("<B", 32))
+        # ... and reserve for the actual SHA256.
+        sha256_loc = self._file.tell()
+        self._file.write(struct.pack("<32B", *[0] * 32))
+        hash_end = self._file.tell()
+        hashes_sz = hash_end - hash_loc - 2
+
+        # Reserve room for our 64-bit tensor length
+        tensor_size_loc = self._file.tell()
+        self._file.write(struct.pack("<Q", 0))
+        tensor_startpos = self._file.tell()
+
+        # Write the total number of bytes for our hash data structures.
+        self._file.seek(hash_loc)
+        self._file.write(struct.pack("<H", hashes_sz))
+        self._file.seek(tensor_startpos)
+
+        tensor_raw_sz = 0
+        tensor_compressed_sz = 0
+        compression_ratio = 0
+        if self.compress_tensors:
+            # NOTE: This compression feature is not complete, as we do not
+            #       yet decompress. This was judged to *not* be worthwhile.
+            #       This is left here as an example for future adventurers that
+            #       may want to do model compression.
+            # Create a write buffer to compress our tensor serialization.
+            tensor_buffer = tempfile.TemporaryFile()
+            tensor.tofile(tensor_buffer)
+            tensor_raw_sz = tensor_buffer.tell()
+            self.total_tensor_bytes += tensor_raw_sz
+            tensor_buffer.seek(0)
+            tensor_compressed = self.lz4_frame.compress(tensor_buffer.read())
+            tensor_compressed_sz = len(tensor_compressed)
+            compression_ratio = (tensor_raw_sz * 1.0) / tensor_compressed_sz
+            if compression_ratio > 2:
+                self.total_compressed_tensor_bytes += tensor_compressed_sz
+            else:
+                self.total_compressed_tensor_bytes += tensor_raw_sz
+
+        # Serialize our tensors
+        tensor.tofile(self._file)
+        tensor_endpos = self._file.tell()
+
+        # Go back and write our tensor length out
+        self._file.seek(tensor_size_loc, io.SEEK_SET)
+        # We write this signed, so that we can use the signedness as an
+        # indicator of possible tensor compression in the future.
+        tensor_size = tensor_endpos - tensor_startpos
+        self._file.write(struct.pack("<q", tensor_size))
+        self.total_tensor_bytes += tensor_size
+
+        # Write our data structure header size.
+        self._file.seek(ds_header_begin)
+        ds_header_size = tensor_startpos - ds_header_begin
+        self._file.write(struct.pack("<Q", ds_header_size))
+
+        # Add our tensor metadata to the index.
+        self._file.seek(self._metadata_cur)
+        self._file.write(struct.pack("<H", len(name_bytes)))
+        self._file.write(name_bytes)
+        self._file.write(struct.pack("<B", tensor_type.value))
+        self._file.write(struct.pack("<B", dtype_len))
+        self._file.write(dtype_bytes)
+        self._file.write(shape_bytes)
+        self._file.write(struct.pack("<Q", ds_header_begin))
+        self._file.write(struct.pack("<Q", tensor_startpos))
+        self._file.write(struct.pack("<Q", tensor_endpos - tensor_startpos))
+        self._metadata_cur = self._file.tell()
+
+        # Check for overflow
+        if self._file.tell() > self._metadata_end:
+            raise RuntimeError("Metadata overflow")
+
+        # Read our header and tensor back in to calculate the hashes.
+        self._file.seek(ds_header_begin)
+        bytes_to_hash = self._file.read(ds_header_size)
+        bytes_to_hash += self._file.read(tensor_endpos - tensor_startpos)
+
+        # Write our hashes out.
+        crc32 = zlib.crc32(bytes_to_hash)
+        self._file.seek(crc32_loc)
+        self._file.write(struct.pack("<I", crc32))
+
+        sha256 = hashlib.sha256(bytes_to_hash).digest()
+        self._file.seek(sha256_loc)
+        self._file.write(sha256)
+
+        # Move to the end of our serialized tensor to prepare for the next one.
+        self._file.seek(tensor_endpos)
+        self._tensors += 1
+
+        # Update our prolog and epilog.
+        self._sync_prologue_state()
+
+        ds_size = self._file.tell() - ds_header_begin
+        ds_bytes = f"{ds_size:,} bytes"
+
+        typ = {
+            TensorType.PARAM: "p",
+            TensorType.BUFFER: "b",
+            TensorType.STATE_DICT: "sd",
+        }[tensor_type]
+
+        if self.compress_tensors:
+            comp_report = (
+                f" - tensor:[raw: {tensor_raw_sz},"
+                + f" compressed: {tensor_compressed_sz},"
+                + f" ratio: {compression_ratio:.2f}]"
+            )
+        else:
+            comp_report = ""
+        logger.info(
+            f"{idx}:{typ}:{name} - {dtype_bytes.decode('utf-8')} - "
+            f"{tensor.shape} -> {ds_bytes}{comp_report}"
+        )
+
+    def write_module(
+        self, m: torch.nn.Module, remove_tensors: bool = False
+    ) -> None:
+        for module_name, module in m.named_modules():
+            for name, param in module.named_parameters(recurse=False):
+                label = module_name + "." + name
+                self.write_tensor(self._idx, label, TensorType.PARAM, param)
+                if remove_tensors:
+                    setattr(module, name, None)
+            for name, buf in module.named_buffers(recurse=False):
+                label = module_name + "." + name
+                self.write_tensor(self._idx, label, TensorType.BUFFER, buf)
+                if remove_tensors:
+                    setattr(module, name, None)
+            self._idx += 1
+
+    def write_state_dict(self, state_dict: Dict):
+        """
+        Write the state_dict to the file in Tensorizer format.
+
+        It is strongly recommended that you use write_module instead of
+        this function, as it will also write out the parameter type,
+        allowing for zero-copy loading of the module with
+        TensorDeserializer.load_into_module.
+        """
+        idx = 0
+        for name, param in state_dict.items():
+            self.write_tensor(idx, name, TensorType.STATE_DICT, param)
diff --git a/tensorizer/stream_io.py b/tensorizer/stream_io.py
new file mode 100644
index 00000000..d96ed295
--- /dev/null
+++ b/tensorizer/stream_io.py
@@ -0,0 +1,571 @@
+import functools
+import io
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+import typing
+import weakref
+from urllib.parse import urlparse
+
+import boto3
+from io import SEEK_SET, SEEK_CUR, SEEK_END
+from typing import Union, Optional, Dict, Any, Tuple
+import shutil
+
+import tensorizer._wide_pipes as _wide_pipes
+
+__all__ = ["open_stream", "CURLStreamFile"]
+
+logger = logging.getLogger(__name__)
+
+curl_path = shutil.which("curl")
+
+default_s3_read_endpoint = "accel-object.ord1.coreweave.com"
+default_s3_write_endpoint = "object.ord1.coreweave.com"
+
+if sys.platform != "win32":
+    _s3_default_config_paths = (os.path.expanduser("~/.s3cfg"),)
+else:
+    # s3cmd generates its config at a different path on Windows by default,
+    # but it may have been manually placed at ~\.s3cfg instead, so check both.
+    _s3_default_config_paths = tuple(
+        map(os.path.expanduser, (r"~\.s3cfg", r"~\AppData\Roaming\s3cmd.ini"))
+    )
+
+
+class _ParsedCredentials(typing.NamedTuple):
+    config_file: Optional[str]
+    s3_endpoint: Optional[str]
+    s3_access_key: Optional[str]
+    s3_secret_key: Optional[str]
+
+
+@functools.lru_cache(maxsize=None)
+def _get_s3cfg_values(
+    config_paths: Optional[
+        Union[
+            Tuple[Union[str, bytes, os.PathLike], ...], str, bytes, os.PathLike
+        ]
+    ] = None
+) -> _ParsedCredentials:
+    """
+    Gets S3 credentials from the .s3cfg file.
+
+    :param config_paths: The sequence of potential file paths to check
+        for s3cmd config settings. If not provided or an empty tuple,
+        platform-specific default search locations are used.
+        When specifying a sequence, this argument must be a tuple,
+        because this function is cached, and that requires
+        all arguments to be hashable.
+    :return: A 4-tuple, config_file, s3_endpoint, s3_access_key, s3_secret_key,
+        where each element may be None if not found,
+        and config_file is the config file path used.
+        If config_file is None, no valid config file
+        was found, and nothing was parsed.
+
+    If the config_paths argument is not provided or is an empty tuple,
+    platform-specific default search locations are used.
+    This function is cached, and hence config_paths must be a (hashable) tuple
+    when specifying a sequence
+    """
+    if not config_paths:
+        config_paths = _s3_default_config_paths
+    elif isinstance(config_paths, (str, bytes, os.PathLike)):
+        config_paths = (config_paths,)
+
+    import configparser
+
+    config = configparser.ConfigParser()
+
+    # Stop on the first path that can be successfully read
+    for config_path in config_paths:
+        if config.read((config_path,)):
+            break
+    else:
+        return _ParsedCredentials(None, None, None, None)
+
+    if "default" not in config:
+        raise ValueError(f"No default section in {config_path}")
+
+    return _ParsedCredentials(
+        config_file=os.fsdecode(config_path),
+        s3_endpoint=config["default"].get("host_base"),
+        s3_access_key=config["default"].get("access_key"),
+        s3_secret_key=config["default"].get("secret_key"),
+    )
+
+
+class CURLStreamFile:
+    """
+    CURLStreamFile implements a file-like object around an HTTP download, the
+    intention being to not buffer more than we have to. It is intended for
+    tar-like files, where we start at the beginning and read until the end of
+    the file.
+
+    It does implement `seek` and `tell`, but only for the purpose of
+    implementing `read`, and only for the purpose of reading the entire file.
+    It does support seeking to an arbitrary position, but is very inefficient
+    in doing so as it requires re-opening the connection to the server.
+    """
+
+    def __init__(
+        self,
+        uri: str,
+        begin: Optional[int] = None,
+        end: Optional[int] = None,
+        headers: Dict[str, Any] = None,
+    ) -> None:
+        self._uri = uri
+
+        if curl_path is None:
+            RuntimeError(
+                "cURL is a required dependency for streaming downloads"
+                " and could not be found."
+            )
+
+        # NOTE: `256mb` buffer on the python IO object.
+        cmd = [
+            curl_path,
+            "--header",
+            "Accept-Encoding: identity",
+            "-s",
+            uri,
+        ]
+
+        if begin is not None or end is not None:
+            cmd.extend(["--range", f"{begin or 0}-{end or ''}"])
+
+        if headers is not None:
+            for k, v in headers.items():
+                cmd.extend(["--header", f"{k}: {v}"])
+
+        self._curl = None
+        with _wide_pipes.widen_new_pipes():  # Widen on Windows
+            self._curl = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                bufsize=256 * 1024 * 1024,
+            )
+
+        _wide_pipes.widen_pipe(self._curl.stdout.fileno())  # Widen on Linux
+        self._curr = 0 if begin is None else begin
+        self._end = end
+        self.closed = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def __del__(self):
+        self.close()
+
+    def _read_until(
+        self, goal_position: int, ba: Union[bytearray, None] = None
+    ) -> Union[bytes, int]:
+        if ba is None:
+            rq_sz = goal_position - self._curr
+            if self._end is not None and self._curr + rq_sz > self._end:
+                rq_sz = self._end - self._curr
+                if rq_sz <= 0:
+                    return bytes()
+            ret_buff = self._curl.stdout.read(rq_sz)
+            ret_buff_sz = len(ret_buff)
+        else:
+            rq_sz = len(ba)
+            if self._end is not None and self._curr + rq_sz > self._end:
+                rq_sz = self._end - self._curr
+                if rq_sz <= 0:
+                    return 0
+                tmp_ba = bytearray(rq_sz)
+                ret_buff_sz = self._curl.stdout.readinto(tmp_ba)
+                ba[:ret_buff_sz] = tmp_ba[:ret_buff_sz]
+                ret_buff = ba
+            else:
+                ret_buff_sz = self._curl.stdout.readinto(ba)
+                ret_buff = ba
+        if ret_buff_sz != rq_sz:
+            self.closed = True
+            self._curl.terminate()
+            if self._curl.returncode != 0:
+                raise IOError(f"curl error: {self._curl.returncode}")
+            else:
+                raise IOError(f"Requested {rq_sz} != {ret_buff_sz}")
+        self._curr += ret_buff_sz
+        if ba is None:
+            return ret_buff
+        else:
+            return ret_buff_sz
+
+    def tell(self) -> int:
+        return self._curr
+
+    def readinto(self, ba: bytearray) -> int:
+        goal_position = self._curr + len(ba)
+        return self._read_until(goal_position, ba)
+
+    def read(self, size=None) -> bytes:
+        if self.closed:
+            raise IOError("CURLStreamFile closed.")
+        if size is None:
+            return self._curl.stdout.read()
+        goal_position = self._curr + size
+        return self._read_until(goal_position)
+
+    @staticmethod
+    def writable() -> bool:
+        return False
+
+    @staticmethod
+    def fileno() -> int:
+        return -1
+
+    def close(self):
+        self.closed = True
+        if self._curl is not None:
+            if self._curl.poll() is None:
+                self._curl.stdout.close()
+                self._curl.terminate()
+                self._curl.wait()
+            else:
+                # stdout is normally closed by the Popen.communicate() method,
+                # which we skip in favour of Popen.stdout.read()
+                self._curl.stdout.close()
+            self._curl = None
+
+    def readline(self):
+        raise NotImplementedError("Unimplemented")
+
+    """
+    This seek() implementation should be avoided if you're seeking backwards,
+    as it's not very efficient due to the need to restart the curl process.
+    """
+
+    def seek(self, position, whence=SEEK_SET):
+        if whence == SEEK_CUR:
+            position += self._curr
+        if position == self._curr:
+            return
+        if whence == SEEK_END:
+            raise ValueError("Unsupported `whence`")
+        elif position > self._curr:
+            # We're seeking forward, so we just read until we get there.
+            self._read_until(position)
+        else:
+            # To seek backwards, we need to close out our existing process and
+            # start a new one.
+            self.close()
+
+            # And we reinitialize ourself.
+            self.__init__(self._uri, position, None)
+
+
+def _ensure_https_endpoint(endpoint: str):
+    scheme, *location = endpoint.split("://", maxsplit=1)
+    scheme = scheme.lower() if location else None
+    if scheme is None:
+        return "https://" + endpoint
+    elif scheme == "https":
+        return endpoint
+    else:
+        raise ValueError("Non-HTTPS endpoint URLs are not allowed.")
+
+
+def s3_upload(
+    path: str,
+    target_uri: str,
+    s3_access_key_id: str,
+    s3_secret_access_key: str,
+    s3_endpoint: str = default_s3_write_endpoint,
+):
+    if s3_secret_access_key is None:
+        raise TypeError("No secret key provided")
+    if s3_access_key_id is None:
+        raise TypeError("No access key provided")
+    if s3_endpoint is None:
+        raise TypeError("No S3 endpoint provided")
+
+    client = boto3.session.Session.client(
+        boto3.session.Session(),
+        endpoint_url=_ensure_https_endpoint(s3_endpoint),
+        service_name="s3",
+        aws_access_key_id=s3_access_key_id,
+        aws_secret_access_key=s3_secret_access_key,
+    )
+    path_uri = urlparse(target_uri)
+    bucket = path_uri.netloc
+    key = path_uri.path.lstrip("/")
+
+    client.upload_file(path, bucket, key)
+
+
+def s3_download(
+    path_uri: str,
+    s3_access_key_id: str,
+    s3_secret_access_key: str,
+    s3_endpoint: str = default_s3_read_endpoint,
+) -> CURLStreamFile:
+    if s3_secret_access_key is None:
+        raise TypeError("No secret key provided")
+    if s3_access_key_id is None:
+        raise TypeError("No access key provided")
+    if s3_endpoint is None:
+        raise TypeError("No S3 endpoint provided")
+
+    client = boto3.session.Session.client(
+        boto3.session.Session(),
+        endpoint_url=_ensure_https_endpoint(s3_endpoint),
+        service_name="s3",
+        aws_access_key_id=s3_access_key_id,
+        aws_secret_access_key=s3_secret_access_key,
+    )
+    path_uri = urlparse(path_uri)
+    bucket = path_uri.netloc
+    key = path_uri.path.lstrip("/")
+
+    url = client.generate_presigned_url(
+        ClientMethod="get_object",
+        Params={"Bucket": bucket, "Key": key},
+        ExpiresIn=300,
+    )
+    return CURLStreamFile(url)
+
+
+def _infer_credentials(
+    s3_access_key_id: Optional[str],
+    s3_secret_access_key: Optional[str],
+    s3_config_path: Optional[Union[str, bytes, os.PathLike]] = None,
+) -> _ParsedCredentials:
+    """
+    Fill in a potentially incomplete S3 credential pair
+    by parsing the s3cmd config file if necessary.
+    An empty string ("") is considered a specified credential,
+    while None is an unspecified credential.
+    Use "" for public buckets.
+
+    :param s3_access_key_id: `s3_access_key_id` if explicitly specified,
+        otherwise None. If None, the s3cmd config file is parsed for the key.
+    :param s3_secret_access_key: `s3_secret_access_key` if explicitly specified,
+        otherwise None. If None, the s3cmd config file is parsed for the key.
+    :param s3_config_path: An explicit path to the s3cmd config file to search,
+        if necessary. If None, platform-specific default paths are used.
+    :return: A `_ParsedCredentials` object with both the
+        `s3_access_key` and `s3_secret_key` fields guaranteed to not be None.
+    :raises ValueError: If the credential pair is incomplete and the
+        missing parts could not be found in any s3cmd config file.
+    """
+    if None not in (s3_access_key_id, s3_secret_access_key):
+        # All required credentials were specified; don't parse anything
+        return _ParsedCredentials(
+            config_file=None,
+            s3_endpoint=None,
+            s3_access_key=s3_access_key_id,
+            s3_secret_key=s3_secret_access_key,
+        )
+
+    # Try to find default credentials if at least one is not specified
+    if s3_config_path is not None and not os.path.exists(s3_config_path):
+        raise FileNotFoundError(
+            f"Explicitly specified s3_config_path does not exist: {s3_config_path}"
+        )
+    try:
+        parsed: _ParsedCredentials = _get_s3cfg_values(s3_config_path)
+    except ValueError as parse_error:
+        raise ValueError(
+            "Attempted to access an S3 bucket,"
+            " but credentials were not provided,"
+            " and the fallback .s3cfg file could not be parsed."
+        ) from parse_error
+
+    if parsed.config_file is None:
+        raise ValueError(
+            "Attempted to access an S3 bucket,"
+            " but credentials were not provided,"
+            " and no default .s3cfg file could be found."
+        )
+
+    # Don't override a specified credential
+    if s3_access_key_id is None:
+        s3_access_key_id = parsed.s3_access_key
+    if s3_secret_access_key is None:
+        s3_secret_access_key = parsed.s3_secret_key
+
+    # Verify that both keys were ultimately found
+    for required_credential, credential_name in (
+        (s3_access_key_id, "s3_access_key_id"),
+        (s3_secret_access_key, "s3_secret_access_key"),
+    ):
+        if not required_credential:
+            raise ValueError(
+                "Attempted to access an S3 bucket,"
+                f" but {credential_name} was not provided,"
+                " and could not be found in the default"
+                f" config file at {parsed.config_file}."
+            )
+
+    return _ParsedCredentials(
+        config_file=parsed.config_file,
+        s3_endpoint=parsed.s3_endpoint,
+        s3_access_key=s3_access_key_id,
+        s3_secret_key=s3_secret_access_key,
+    )
+
+
+def _temp_file_closer(file: io.IOBase, file_name: str, *upload_args):
+    """
+    Close, upload by name, and then delete the file.
+    Meant to replace .close() on a particular instance
+    of a temporary file-like wrapper object, as an unbound
+    callback to a weakref.finalize() registration on the wrapper.
+
+    The reason this implementation is necessary is really complicated.
+
+    ---
+
+    boto3's upload_fileobj could be used before closing the
+    file, instead of closing it and then uploading it by
+    name, but upload_fileobj is less performant than
+    upload_file as of boto3's s3 library s3transfer
+    version 0.6.0.
+
+    For details, see the implementation & comments:
+    https://github.com/boto/s3transfer/blob/0.6.0/s3transfer/upload.py#L351
+
+    TL;DR: s3transfer does multithreaded transfers
+    that require multiple file handles to work properly,
+    but Python cannot duplicate file handles such that
+    they can be accessed in a thread-safe way,
+    so they have to buffer it all in memory.
+    """
+
+    if file.closed:
+        # Makes closure idempotent.
+
+        # If the file object is used as a context
+        # manager, close() is called twice (once in the
+        # serializer code, once after, when leaving the
+        # context).
+
+        # Without this check, this would trigger two
+        # separate uploads.
+        return
+    try:
+        file.close()
+        s3_upload(file_name, *upload_args)
+    finally:
+        try:
+            os.unlink(file_name)
+        except OSError:
+            pass
+
+
+def open_stream(
+    path_uri: Union[str, os.PathLike],
+    mode: str = "rb",
+    s3_access_key_id: Optional[str] = None,
+    s3_secret_access_key: Optional[str] = None,
+    s3_endpoint: Optional[str] = None,
+    s3_config_path: Optional[Union[str, bytes, os.PathLike]] = None,
+) -> Union[CURLStreamFile, typing.BinaryIO]:
+    """Open a file path, http(s):// URL, or s3:// URI.
+    :param path_uri: File path, http(s):// URL, or s3:// URI to open.
+    :param mode: Mode with which to open the stream.
+        Supported values are:
+        * "rb" for http(s)://,
+        * "rb", "wb[+]", and "ab[+]" for s3://,
+        * All standard binary modes for file paths.
+    :param s3_access_key_id: S3 access key, corresponding to
+        "aws_access_key_id" in boto3.
+        If not specified, an s3:// URI is being opened, and ~/.s3cfg exists,
+        ~/.s3cfg's "access_key" will be parsed as this credential.
+        To specify blank credentials, for a public bucket,
+        pass the empty string ("") rather than None.
+    :param s3_secret_access_key: S3 secret key, corresponding to
+        "aws_secret_access_key" in boto3.
+        If not specified, an s3:// URI is being opened, and ~/.s3cfg exists,
+        ~/.s3cfg's "secret_key" will be parsed as this credential.
+        To specify blank credentials, for a public bucket,
+        pass the empty string ("") rather than None.
+    :param s3_endpoint: S3 endpoint.
+        If not specified and a host_base was found
+        alongside previously parsed credentials, that will be used.
+        Otherwise, object.ord1.coreweave.com is the default.
+    :param s3_config_path: An explicit path to the ~/.s3cfg config file
+        to be parsed if full credentials are not provided.
+        If None, platform-specific default paths are used.
+    :return: An opened file-like object representing the target resource.
+    """
+    if isinstance(path_uri, os.PathLike):
+        path_uri = os.fspath(path_uri)
+
+    scheme, *location = path_uri.split("://", maxsplit=1)
+    scheme = scheme.lower() if location else None
+
+    normalized_mode = "".join(sorted(mode))
+
+    if scheme in ("http", "https"):
+        if normalized_mode != "br":
+            raise ValueError(
+                'Only the mode "rb" is valid when opening http(s):// streams.'
+            )
+        return CURLStreamFile(path_uri)
+
+    elif scheme == "s3":
+        if normalized_mode not in ("br", "bw", "ab", "+bw", "+ab"):
+            raise ValueError(
+                'Only the modes "rb", "wb[+]", and "ab[+]" are valid'
+                " when opening s3:// streams."
+            )
+
+        s3 = _infer_credentials(
+            s3_access_key_id, s3_secret_access_key, s3_config_path
+        )
+        s3_access_key_id = s3.s3_access_key
+        s3_secret_access_key = s3.s3_secret_key
+
+        # Not required to have been found,
+        # and doesn't overwrite an explicitly specified endpoint.
+        s3_endpoint = s3_endpoint or s3.s3_endpoint
+
+        # Regardless of whether the config needed to be parsed,
+        # the endpoint gets a default value.
+
+        if "w" in mode or "a" in mode:
+            s3_endpoint = s3_endpoint or default_s3_write_endpoint
+
+            # delete must be False or the file will be deleted by the OS
+            # as soon as it closes, before it can be uploaded on platforms
+            # with primitive temporary file support (e.g. Windows)
+            temp_file = tempfile.NamedTemporaryFile(mode="wb+", delete=False)
+
+            guaranteed_closer = weakref.finalize(
+                temp_file,
+                _temp_file_closer,
+                temp_file.file,
+                temp_file.name,
+                path_uri,
+                s3_access_key_id,
+                s3_secret_access_key,
+                s3_endpoint,
+            )
+            temp_file.close = guaranteed_closer
+            return temp_file
+        else:
+            s3_endpoint = s3_endpoint or default_s3_read_endpoint
+            return s3_download(
+                path_uri, s3_access_key_id, s3_secret_access_key, s3_endpoint
+            )
+
+    else:
+        if "b" not in normalized_mode:
+            raise ValueError(
+                'Only binary modes ("rb", "wb", "wb+", etc.)'
+                " are valid when opening local file streams."
+            )
+        os.makedirs(os.path.dirname(path_uri), exist_ok=True)
+        handle: typing.BinaryIO = open(path_uri, mode)
+        handle.seek(0)
+        return handle
diff --git a/tensorizer/tensorizer.py b/tensorizer/tensorizer.py
deleted file mode 100644
index 8d1699b5..00000000
--- a/tensorizer/tensorizer.py
+++ /dev/null
@@ -1,1015 +0,0 @@
-##############################################################################
-# tensorizer.py                                                      Wes Brown
-# Fast LLM model serialization and deserialization          (c) 2022 Coreweave
-##############################################################################
-
-# try to import UNIX only dependencies
-try:
-    import fcntl
-    import resource
-except ImportError:
-    fcntl = None
-    resource = None
-
-import subprocess
-import io
-from io import SEEK_SET, SEEK_END
-import numpy
-import struct
-import torch
-import typing
-import time
-import logging
-import sys
-import pathlib
-import os
-import requests
-import argparse
-
-os.environ[
-    "TRANSFORMERS_VERBOSITY"
-] = "error"  # disable missing keys and unexpected key warnings
-
-thisPath = str(pathlib.Path(__file__).parent.resolve())
-sys.path.append(thisPath + "/../../transformers/src")
-sys.path.append(thisPath + "/../../")
-
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    AutoConfig,
-    CLIPTextModel,
-    CLIPTokenizer,
-    CLIPTextConfig,
-)
-from transformers.modeling_utils import no_init_weights, PreTrainedModel
-from tokenizers import Tokenizer
-
-from diffusers import (
-    AutoencoderKL,
-    UNet2DConditionModel,
-    StableDiffusionPipeline,
-    LMSDiscreteScheduler,
-)
-from diffusers.modeling_utils import ModelMixin
-from diffusers.configuration_utils import ConfigMixin
-
-import json
-from urllib import request
-import tempfile
-
-from collections import OrderedDict
-from typing import Optional, Tuple, Union, List, Iterator, Callable
-
-lz4 = None
-
-# Setup logger
-logger = logging.getLogger(__name__)
-logger.setLevel(level=logging.INFO)
-fh = logging.StreamHandler()
-fh_formatter = logging.Formatter(
-    "%(asctime)s %(levelname)s %(filename)s(%(process)d) - %(message)s"
-)
-fh.setFormatter(fh_formatter)
-logger.addHandler(fh)
-
-# =============================================================================
-# From `pipe(7)` manpage:
-#
-# Pipe capacity
-# A pipe has a limited capacity. If the pipe is full, then a write(2) will
-# block or fail, depending on whether the O_NONBLOCK flag is set (see below).
-# Different implementations have different limits for the pipe capacity.
-#
-# Applications should not rely on a particular capacity: an application should
-# be designed so that a reading process consumes data as soon as it is
-# available, so that a writing process does not remain blocked.
-#
-# In Linux versions before 2.6.11, the capacity of a pipe was the same as the
-# system page size (e.g., 4096 bytes on i386). Since Linux 2.6.11, the pipe
-# capacity is 16 pages (i.e., 65,536 bytes in a system with a page size of
-# 4096 bytes). Since Linux 2.6.35, the default pipe capacity is 16 pages, but
-# the capacity can be queried and set using the fcntl(2) F_GETPIPE_SZ and
-# F_SETPIPE_SZ operations. See fcntl(2) for more information.
-#
-# =============================================================================
-# From `fcntl(2)` manpage:
-#
-# Changing the capacity of a pipe
-#
-# F_SETPIPE_SZ (int; since Linux 2.6.35)
-# Change the capacity of the pipe referred to by fd to be at least arg bytes.
-# An unprivileged process can adjust the pipe capacity to any value between the
-# system page size and the limit defined in /proc/sys/fs/pipe−max−size
-# (see proc(5)). Attempts to set the pipe capacity below the page size are
-# silently rounded up to the page size. Attempts by an unprivileged process to
-# set the pipe capacity above the limit in /proc/sys/fs/pipe−max−size yield the
-# error EPERM; a privileged process (CAP_SYS_RESOURCE) can override the limit.
-#
-# When allocating the buffer for the pipe, the kernel may use a capacity larger
-# than arg, if that is convenient for the implementation. (In the current
-# implementation, the allocation is the next higher power-of-two page-size
-# multiple of the requested size.) The actual capacity (in bytes) that is set
-# is returned as the function result.
-#
-# Attempting to set the pipe capacity smaller than the amount of buffer space
-# currently used to store data produces the error EBUSY.
-#
-# Note that because of the way the pages of the pipe buffer are employed when
-# data is written to the pipe, the number of bytes that can be written may be
-# less than the nominal size, depending on the size of the writes.
-#
-# F_GETPIPE_SZ (void; since Linux 2.6.35)
-# Return (as the function result) the capacity of the pipe referred to by fd.
-#
-# =============================================================================
-# Constant for `F_SETPIPE_SZ`, as python3's `fcntl` module doesn't have this
-# defined -- despite the documentation saying that they're there.
-#
-# TODO: Make this work or fail gracefully on non-Linux systems. Not sure if
-#       this is really relevant, as I don't even know if CUDA is available on
-#       non-Linux systems in a production sense.
-F_SETPIPE_SZ = 1031
-
-# Whether the tensor is a parameter or a buffer on the model.
-TENSOR_PARAM = 0
-TENSOR_BUFFER = 1
-
-
-# Silly function to convert to human bytes
-def convert_bytes(num):
-    """
-    this function will convert bytes to MB.... GB... etc
-    """
-    step_unit = 1000.0
-
-    for x in ["bytes", "KB", "MB", "GB", "TB"]:
-        if num < step_unit:
-            return "%3.1f %s" % (num, x)
-        num /= step_unit
-
-
-def no_init_or_tensor(loading_code):
-    def dummy(self):
-        return
-
-    modules = [torch.nn.Linear, torch.nn.Embedding, torch.nn.LayerNorm]
-    original = {}
-    for mod in modules:
-        original[mod] = mod.reset_parameters
-        mod.reset_parameters = dummy
-    original_empty = torch.empty
-
-    # when torch.empty is called, make it map to meta device by replacing the device in kwargs.
-    torch.empty = lambda *args, **kwargs: original_empty(
-        *args, **{**kwargs, "device": "meta"}
-    )
-
-    with no_init_weights():
-        result = loading_code()
-    for mod in modules:
-        mod.reset_parameters = original[mod]
-    torch.empty = original_empty
-
-    return result
-
-
-class CURLStreamFile(object):
-    """
-    CURLStreamFile implements a file-like object around an HTTP download, the
-    intention being to not buffer more than we have to.
-    """
-
-    def __init__(self, uri: str) -> None:
-        # NOTE: `256mb` buffer on the python IO object.
-        self._curl = subprocess.Popen(
-            [
-                "/usr/bin/curl",
-                "--header",
-                "Accept-Encoding: identity",
-                "-s",
-                uri,
-            ],
-            stdout=subprocess.PIPE,
-            bufsize=256 * 1024 * 1024,
-        )
-        # Read our max-fd-size, fall back to 1mb if invalid.
-        pipe_buf_sz = 1024 * 1024
-        try:
-            pipe_file = open("/proc/sys/fs/pipe-max-size", "r")
-            pipe_buf_sz = int(pipe_file.read())
-            logger.debug(f"pipe-max-size: {pipe_buf_sz}")
-        except IOError as e:
-            logger.warning(
-                f"Could not read /proc/sys/fs/pipe-max-size: {e.strerror}"
-            )
-        try:
-            fcntl.fcntl(self._curl.stdout.fileno(), F_SETPIPE_SZ, pipe_buf_sz)
-        except PermissionError as e:
-            logger.warning(
-                f"Couldn't fcntl F_SETPIPE_SZ to {pipe_buf_sz}: {e.strerror}"
-            )
-        self._curr = 0
-        self.closed = False
-
-    def _read_until(
-        self, goal_position: int, ba: Union[bytearray, None] = None
-    ) -> Union[bytes, int]:
-        if ba is None:
-            rq_sz = goal_position - self._curr
-            ret_buff = self._curl.stdout.read(rq_sz)
-            ret_buff_sz = len(ret_buff)
-        else:
-            rq_sz = len(ba)
-            ret_buff_sz = self._curl.stdout.readinto(ba)
-            ret_buff = ba
-        if ret_buff_sz != rq_sz:
-            self.closed = True
-            err = self._curl.stderr.read()
-            self._curl.terminate()
-            if self._curl.returncode != 0:
-                raise (IOError(f"curl error: {self._curl.returncode}, {err}"))
-            else:
-                raise (IOError(f"Requested {rq_sz} != {ret_buff_sz}"))
-        self._curr += ret_buff_sz
-        if ba is None:
-            return ret_buff
-        else:
-            return ret_buff_sz
-
-    def tell(self) -> int:
-        return self._curr
-
-    def readinto(self, ba: bytearray) -> int:
-        goal_position = self._curr + len(ba)
-        return self._read_until(goal_position, ba)
-
-    def read(self, size=None) -> bytes:
-        if self.closed:
-            raise (IOError("CURLStreamFile closed."))
-        if size is None:
-            return self._curl.stdout.read()
-        goal_position = self._curr + size
-        return self._read_until(goal_position)
-
-    @staticmethod
-    def writable() -> bool:
-        return False
-
-    @staticmethod
-    def fileno() -> int:
-        return -1
-
-    def close(self):
-        self.closed = True
-        self._curl.terminate()
-
-    def readline(self):
-        raise Exception("Unimplemented")
-
-    """
-    This seek() implementation is effectively a no-op, and will throw an
-    exception for anything other than a seek to the current position.
-    """
-
-    def seek(self, position, whence=SEEK_SET):
-        if position == self._curr:
-            return
-        if whence == SEEK_END:
-            raise (Exception("Unsupported `whence`"))
-        else:
-            raise (Exception("Seeking is unsupported"))
-
-
-class RequestsStreamFile(object):
-    """
-    RequestsStreamFile implements a file-like object around an HTTP download, the
-    intention being to not buffer more than we have to. Not as fast or efficient
-    as CURLStreamFile, but it works on Windows.
-    """
-
-    def __init__(self, uri: str) -> None:
-        self._uri = uri
-        self._curr = 0
-        self._r = requests.get(uri, stream=True)
-        self.closed = False
-
-    def _read_until(
-        self, goal_position: int, ba: Union[bytearray, None] = None
-    ) -> Union[bytes, int]:
-        if ba is None:
-            rq_sz = goal_position - self._curr
-            ret_buff = self._r.raw.read(rq_sz)
-            ret_buff_sz = len(ret_buff)
-        else:
-            rq_sz = len(ba)
-            ret_buff_sz = self._r.raw.readinto(ba)
-            ret_buff = ba
-        if ret_buff_sz != rq_sz:
-            self.closed = True
-            raise (IOError(f"Requested {rq_sz} != {ret_buff_sz}"))
-        self._curr += ret_buff_sz
-        if ba is None:
-            return ret_buff
-        else:
-            return ret_buff_sz
-
-    def tell(self) -> int:
-        return self._curr
-
-    def readinto(self, ba: bytearray) -> int:
-        goal_position = self._curr + len(ba)
-        return self._read_until(goal_position, ba)
-
-    def read(self, size=None) -> bytes:
-        if self.closed:
-            raise (IOError("RequestsStreamFile closed."))
-        if size is None:
-            return self._r.raw.read()
-        goal_position = self._curr + size
-        return self._read_until(goal_position)
-
-    @staticmethod
-    def writable() -> bool:
-        return False
-
-    @staticmethod
-    def fileno() -> int:
-        return -1
-
-    def close(self):
-        self.closed = True
-        self._r.close()
-        del self._r
-
-    def readline(self):
-        raise Exception("Unimplemented")
-
-    """
-    This seek() implementation is effectively a no-op, and will throw an
-    exception for anything other than a seek to the current position.
-    """
-
-    def seek(self, position, whence=SEEK_SET):
-        if position == self._curr:
-            return
-        if whence == SEEK_END:
-            raise (Exception("Unsupported `whence`"))
-        else:
-            raise (Exception("Seeking is unsupported"))
-
-
-class GooseTensorizer:
-    """
-    Given a file-like object, either for read or write, serialize tensors
-    to the file, or deserialize the tensors.
-    """
-
-    def __init__(self, file_obj, compress_tensors=False) -> None:
-        self._file = file_obj
-        self._tensors = 0
-        self.total_tensor_bytes = 0
-        self.total_compressed_tensor_bytes = 0
-        self.compress_tensors = compress_tensors
-        self.read_bytes = 0
-        self._buffer = bytearray(1024 * 1024)
-        if self.compress_tensors:
-            import lz4.frame
-
-    @staticmethod
-    def _dump_shape(obj) -> Tuple[bytes, int]:
-        """
-        Returns shape of the tensor
-        """
-        bstr = struct.pack("<B", len(obj))
-        for i in obj:
-            bstr += struct.pack("<I", i)
-        return bstr
-
-    @staticmethod
-    def _read_shapes(obj, num_elems) -> List[int]:
-        """
-        Read the tensor shapes.
-        """
-        bstr = obj
-        shape = []
-        for i in range(num_elems):
-            shape.append(struct.unpack("<I", bstr[0:4])[0])
-            bstr = bstr[4:]
-        return shape
-
-    @property
-    def total_bytes_read(self) -> int:
-        if self._file.closed:
-            return self.total_tensor_bytes
-        else:
-            return self._file.tell()
-
-    def finalize(self) -> None:
-        """
-        Finalizes the serialization with a zero-length field, and closes out
-        the file.
-        """
-        self._file.write(struct.pack("<Q", 0))
-        final_sz = self._file.tell()
-        self._file.close()
-        logger.info(f"Tensors completed serializing to {final_sz} bytes")
-        if self.compress_tensors:
-            compression_ratio = (
-                self.total_tensor_bytes / self.total_compressed_tensor_bytes
-            )
-            logger.info(f"Uncomp'd bytes: {self.total_tensor_bytes}")
-            logger.info(f"Comp'd bytes: {self.total_compressed_tensor_bytes}")
-            logger.info(f"Ratio: {compression_ratio:.2f}")
-
-    def write_tensor(self, idx, name, tensor_type, tensor) -> None:
-        """
-        Serializes a tensor, laying things out so that it can be read in three
-        calls from the input -- once for the size, once for the header, and
-        once for the tensor itself.
-
-        Serialization format:
-
-         { uint64                           header_sz,
-           uint16                           module_idx,
-           uint8                            type,
-           uint16                           name_sz,
-           []char                           name,
-           uint8                            dtype_sz,
-           []char                           dtype_str,
-           []{uint8                         shape_elem_type,
-              union{uint8, uint16, uint32}  shape_elem_sz,
-             }                              shape_elements,
-           uint64                           tensor_sz,
-           []byte                           tensor }
-        """
-
-        if len(str(tensor.dtype)) >= 256:
-            raise ValueError("dtype length should be less than 256")
-
-        # Reserve room for our tensor header size.
-        ds_header_begin = self._file.tell()
-        self._file.write(struct.pack("<Q", 0))
-
-        # Module index.
-        self._file.write(struct.pack("<H", idx))
-
-        # Whether this is a parameter or a buffer
-        self._file.write(struct.pack("<B", tensor_type))
-
-        # Parameter/buffer name
-        name_bytes = bytes(name, "utf-8")
-        self._file.write(struct.pack("<H", len(name_bytes)))
-        self._file.write(name_bytes)
-
-        # Write out our tensor dtype
-        dtype_bytes = bytes(str(tensor.dtype), "utf-8")
-        dtype_len = len(dtype_bytes)
-        self._file.write(struct.pack("<B", dtype_len))
-        self._file.write(dtype_bytes)
-
-        # ... and shape
-        self._file.write(self._dump_shape(tensor.shape))
-
-        # Reserve room for our 64-bit tensor length
-        tensor_size_loc = self._file.tell()
-        self._file.write(struct.pack("<q", 0))
-
-        if self.compress_tensors:
-            # NOTE: This compression feature is not complete, as we do not
-            #       yet decompress. This was judged to *not* be worthwhile.
-            #       This is left here as an example for future adventurers that
-            #       may want to do model compression.
-            # Create a write buffer to compress our tensor serialization.
-            tensor_buffer = tempfile.TemporaryFile()
-            tensor.tofile(tensor_buffer)
-            tensor_raw_sz = tensor_buffer.tell()
-            self.total_tensor_bytes += tensor_raw_sz
-            tensor_buffer.seek(0)
-            tensor_compressed = lz4.frame.compress(tensor_buffer.read())
-            tensor_compressed_sz = len(tensor_compressed)
-            compression_ratio = (tensor_raw_sz * 1.0) / tensor_compressed_sz
-            if compression_ratio > 2:
-                self.total_compressed_tensor_bytes += tensor_compressed_sz
-            else:
-                self.total_compressed_tensor_bytes += tensor_raw_sz
-
-        # Serialize our tensors
-        tensor_startpos = self._file.tell()
-        tensor.tofile(self._file)
-        tensor_endpos = self._file.tell()
-
-        # Go back and write our tensor length out
-        self._file.seek(tensor_size_loc, io.SEEK_SET)
-        # We write this signed, so that we can use the signedness as an
-        # indicator of possible tensor compression in the future.
-        self._file.write(struct.pack("<q", tensor_endpos - tensor_startpos))
-
-        # Write our data structure header size.
-        self._file.seek(ds_header_begin)
-        ds_header_size = tensor_startpos - ds_header_begin
-        self._file.write(struct.pack("<Q", ds_header_size))
-
-        # Move to the end of our serialized tensor to prepare for the next one.
-        self._file.seek(tensor_endpos)
-        self._tensors += 1
-        ds_size = self._file.tell() - ds_header_begin
-        ds_bytes = f"{ds_size:,} bytes"
-
-        if tensor_type == TENSOR_PARAM:
-            typ = "p"
-        else:
-            typ = "b"
-        if self.compress_tensors:
-            comp_report = (
-                f" - tensor:[raw: {tensor_raw_sz},"
-                + f" compressed: {tensor_compressed_sz},"
-                + f" ratio: {compression_ratio:.2f}]"
-            )
-        else:
-            comp_report = ""
-        logger.info(
-            f"{idx}:{typ}:{name} - {tensor.shape} -> {ds_bytes}{comp_report}"
-        )
-
-    def read_tensors(self) -> Iterator[Tuple[int, int, str, numpy.ndarray]]:
-        """
-        A generator that deserializes tensors and returns the `module_idx`,
-        `tensor_type`, parameter/buffer `name`, and the numpy `arr` that
-        represents the tensor.
-        """
-        try:
-            while True:
-                header_sz = struct.unpack("<Q", self._file.read(8))[0]
-                if header_sz == 0:
-                    break
-                headers = self._file.read(header_sz - 8)
-                header_len = len(headers)
-                module_idx = struct.unpack("<H", headers[0:2])[0]
-                tensor_type = struct.unpack("<B", headers[2:3])[0]
-                name_sz = struct.unpack("<H", headers[3:5])[0]
-                idx = name_sz + 5
-                name_bytes = headers[5:idx]
-                name: str = name_bytes.decode("utf-8")
-                dtype_len = struct.unpack("<B", headers[idx : idx + 1])[0]
-                dtype_end = idx + dtype_len + 1
-                dtype = headers[idx + 1 : dtype_end]
-                # read the shape amount, according to the serialized format. the shape length is
-                # 1 byte after the dtype end.
-                shape_len = struct.unpack(
-                    "<B", headers[dtype_end : dtype_end + 1]
-                )[0]
-                # the shape elements are <I so we read 4 bytes. _read_shapes takes in the
-                # header object and the number of elements in the shape.
-                # the amount of bytes for the shape is 4 * the number of elements in the shape.
-                # so, we need to read 4 * shape_len bytes after the dtype end + 1 byte for the
-                # shape length. sort of convoluted, but it works.
-                shapelist = self._read_shapes(
-                    headers[dtype_end + 1 : (dtype_end + 1) + (4 * shape_len)],
-                    shape_len,
-                )
-                datalength = struct.unpack("<q", headers[header_len - 8 :])[0]
-                if datalength > len(self._buffer):
-                    self._buffer = bytearray(datalength)
-                self._file.readinto(memoryview(self._buffer)[:datalength])
-                arr = numpy.ndarray.__new__(
-                    numpy.memmap,
-                    shapelist,
-                    dtype=dtype,
-                    buffer=self._buffer,
-                    offset=0,
-                )
-                # Get rid of the error on load.
-                yield module_idx, tensor_type, name, arr
-        except EOFError:
-            return
-
-    def load_tensors_dict(
-        self, device=torch.device("cuda"), dtype: str = "float16"
-    ) -> (int, OrderedDict):
-        tensor_ct = 0
-        d = OrderedDict()
-        for idx, typ, name, arr in self.read_tensors():
-            gradient = False
-            if arr.dtype != "bool":
-                gradient = True
-                if arr.dtype != dtype:
-                    arr = arr.astype(dtype)
-            d[name] = torch.nn.Parameter(
-                torch.as_tensor(arr, device=device), requires_grad=gradient
-            )
-            tensor_ct += 1
-        self.total_tensor_bytes = self._file.tell()
-        self._file.close()
-        return tensor_ct, d
-
-    def load_tensors(
-        self,
-        m: torch.nn.Module,
-        device=torch.device("cuda"),
-        dtype: [None, str] = None,
-    ) -> int:
-        """
-        Given `m`, a torch.nn.Module, load the associate tensors in this
-        GooseTensorizer object into the `torch.nn.Module`. Returns the number
-        of tensors loaded into the model.
-
-        Note that this assumes that the `m` is serialized along with this.
-        """
-        if self._file.closed:
-            raise IOError("IO closed, instantiate if you want to load again.")
-
-        modules: typing.OrderedDict[str, torch.nn.Module] = OrderedDict()
-        for name, module in m.named_modules():
-            modules[name] = module
-
-        tensor_ct = 0
-        for idx, typ, name, arr in self.read_tensors():
-            gradient = True
-            if arr.dtype not in ["float", "complex"]:
-                gradient = False
-                if dtype is not None and arr.dtype != dtype:
-                    arr = arr.astype(dtype)
-            obj_path, attr = name.rsplit(".", 1)
-            module: torch.nn.Module = modules[obj_path]
-            param = torch.nn.Parameter(
-                torch.as_tensor(arr, device=device), requires_grad=gradient
-            )
-            if typ is TENSOR_PARAM:
-                module.register_parameter(attr, param)
-            elif typ is TENSOR_BUFFER:
-                module.register_buffer(attr, param)
-            tensor_ct += 1
-        self.total_tensor_bytes = self._file.tell()
-        self._file.close()
-        return tensor_ct
-
-
-def save_to_state_dict(m, destination, prefix, keep_vars):
-    r"""Saves module state to `destination` dictionary, containing a state
-    of the module, but not its descendants. This is called on every
-    submodule in :meth:`~torch.nn.Module.state_dict`.
-
-    In rare cases, subclasses can achieve class-specific behavior by
-    overriding this method with custom logic.
-
-    Args:
-        destination (dict): a dict where state will be stored
-        prefix (str): the prefix for parameters and buffers used in this
-            module
-    """
-    for name, param in m._parameters.items():
-        destination[prefix + name] = param if keep_vars else param.detach()
-    for name, buf in m._buffers.items():
-        if name not in m._non_persistent_buffers_set:
-            destination[prefix + name] = buf if keep_vars else buf.detach()
-    extra_state_key = prefix + "_extra_state"
-    if (
-        getattr(m.__class__, "get_extra_state", torch.nn.Module.get_extra_state)
-        is not torch.nn.Module.get_extra_state
-    ):
-        destination[extra_state_key] = m.get_extra_state()
-
-
-def state_dict(
-    m: torch.nn.Module, destination=None, prefix="", keep_vars=True
-) -> OrderedDict:
-    if destination is None:
-        destination = OrderedDict()
-        destination._metadata = OrderedDict()
-    destination._metadata[prefix[:-1]] = local_metadata = dict(
-        version=m._version
-    )
-    save_to_state_dict(m, destination, prefix, keep_vars)
-    # m._save_to_state_dict(destination, prefix, keep_vars)
-    for name, module in m._modules.items():
-        state_dict(
-            module, destination, prefix + name + ".", keep_vars=keep_vars
-        )
-    for hook in m._state_dict_hooks.values():
-        hook_result = hook(m, destination, prefix, local_metadata)
-        if hook_result is not None:
-            destination = hook_result
-    return destination
-
-
-def get_ram_usage_str() -> str:
-    if resource is not None:
-        maxrss_b4 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (
-            1000 * 1000
-        )
-        maxrss_b4_gb = f"{maxrss_b4:0.2f}gb CPU RAM used"
-    else:
-        maxrss_b4_gb = "unknown CPU RAM used"
-    return maxrss_b4_gb
-
-
-def get_vram_usage_str() -> str:
-    if torch.cuda.is_available():
-        gb_gpu = int(
-            torch.cuda.get_device_properties(0).total_memory
-            / (1000 * 1000 * 1000)
-        )
-        return f"{str(gb_gpu)}gb"
-    return "N/A"
-
-
-def get_gpu_name() -> str:
-    if torch.cuda.is_available():
-        return torch.cuda.get_device_name(0)
-    return "N/A"
-
-
-def get_device() -> torch.device:
-    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-def serialize_model(
-    model: torch.nn.Module,
-    config: Optional[Union[ConfigMixin, AutoConfig, dict]],
-    model_directory: str,
-    model_prefix: str = "model",
-):
-    """
-    Remove the tensors from a PyTorch model, convert them to NumPy
-    arrays and serialize them to GooseTensor format. The stripped
-    model is also serialized to pytorch format.
-
-    Args:
-        model: The model to serialize.
-        config: The model's configuration. This is optional and only
-            required for HuggingFace Transformers models. Diffusers
-            models do not require this.
-        model_directory: The directory to save the serialized model to.
-        model_prefix: The prefix to use for the serialized model files. This
-            is purely optional and it allows for multiple models to be
-            serialized to the same directory. A good example are Stable
-            Diffusion models. Default is "model".
-    """
-
-    os.makedirs(model_directory, exist_ok=True)
-    dir_prefix = f"{model_directory}/{model_prefix}"
-
-    if config is None:
-        config = model
-    if config is not None:
-        if hasattr(config, "to_json_file"):
-            config.to_json_file(f"{dir_prefix}-config.json")
-        if isinstance(config, dict):
-            open(f"{dir_prefix}-config.json", "w").write(
-                json.dumps(config, indent=2)
-            )
-
-    ts = GooseTensorizer(open(f"{dir_prefix}.tensors", "wb"))
-
-    idx = 0
-    for module_name, module in model.named_modules():
-        for name, param in module.named_parameters(recurse=False):
-            v = param.cpu().detach().numpy()
-            ts.write_tensor(idx, module_name + "." + name, TENSOR_PARAM, v)
-            setattr(module, name, None)
-        for name, buf in module.named_buffers(recurse=False):
-            v = buf.cpu().detach().numpy()
-            ts.write_tensor(idx, module_name + "." + name, TENSOR_BUFFER, v)
-            setattr(module, name, None)
-        idx += 1
-
-    ts.finalize()
-
-
-def load_model(
-    path_uri: str,
-    modelclass: Union[PreTrainedModel, ModelMixin, ConfigMixin] = None,
-    configclass: Optional[Union[ConfigMixin, AutoConfig]] = None,
-    model_prefix: str = "model",
-    device: torch.device = "cpu",
-    dtype: str = None,
-) -> torch.nn.Module:
-    """
-    Given a path prefix, load the model with a custom extension
-
-    Args:
-        path_uri: path to the model. Can be a local path or a URI
-        modelclass: The model class to load the tensors into.
-        configclass: The config class to load the model config into. This must be
-            set if you are loading a model from HuggingFace Transformers.
-        model_prefix: The prefix to use to distinguish between multiple serialized models. The default is "model".
-        dtype: The dtype to load the tensors into. If None, the dtype is inferred from the model.
-    """
-
-    if model_prefix is None:
-        model_prefix = "model"
-
-    if path_uri.startswith("https://") or path_uri.startswith("http://"):
-        config_uri = f"{path_uri}/{model_prefix}-config.json"
-        tensors_uri = f"{path_uri}/{model_prefix}.tensors"
-        if fcntl is not None:
-            # We have fcntl, so we can use the fast CURL-based loader.
-            logger.info("Using CURL for tensor streaming")
-            tensor_loader = lambda: CURLStreamFile(tensors_uri)
-        else:
-            # Fallback to slow requests-based loader.
-            logger.info("Using requests for tensor streaming")
-            tensor_loader = lambda: RequestsStreamFile(tensors_uri)
-    else:
-        tensors_uri = f"{path_uri}/{model_prefix}.tensors"
-        config_uri = "file://" + os.path.join(
-            os.path.dirname(path_uri), f"{model_prefix}-config.json"
-        )
-        if not os.path.exists(config_uri):
-            config_uri = f"{path_uri}/{model_prefix}-config.json"
-        tensor_loader = lambda: open(tensors_uri, "rb")
-
-    logger.info(f"Loading {tensors_uri}, {get_ram_usage_str()}")
-    begin_load = time.time()
-
-    tensor_deserializer = GooseTensorizer(tensor_loader())
-
-    if configclass is not None:
-        try:
-            with tempfile.TemporaryDirectory() as dir:
-                request.urlretrieve(
-                    config_uri, os.path.join(dir, "config.json")
-                )
-                config = configclass.from_pretrained(dir)
-                config.gradient_checkpointing = True
-        except ValueError:
-            config = configclass.from_pretrained(config_uri)
-        model = no_init_or_tensor(
-            lambda: modelclass.from_pretrained(
-                None, config=config, state_dict=OrderedDict()
-            )
-        )
-    else:
-        try:
-            config = json.loads(
-                request.urlopen(config_uri).read().decode("utf-8")
-            )
-        except ValueError:
-            with open(config_uri, "r") as f:
-                config = json.load(f)
-        model = no_init_or_tensor(lambda: modelclass(**config))
-
-    tensor_deserializer.load_tensors(model, device=device, dtype=dtype)
-
-    tensor_load_s = time.time() - begin_load
-    rate_str = convert_bytes(
-        tensor_deserializer.total_bytes_read / tensor_load_s
-    )
-    tensors_sz = convert_bytes(tensor_deserializer.total_bytes_read)
-    logger.info(
-        f"Model tensors loaded in {tensor_load_s:0.2f}s, read "
-        + f"{tensors_sz} @ {rate_str}/s, {get_ram_usage_str()}"
-    )
-
-    return model
-
-
-def df_main(args: argparse.Namespace) -> None:
-    output_prefix = args.output_prefix
-    print("MODEL PATH:", args.input_directory)
-    print("OUTPUT PREFIX:", output_prefix)
-
-    hf_api_token = os.environ.get("HF_API_TOKEN")
-
-    pipeline = StableDiffusionPipeline.from_pretrained(
-        args.input_directory, use_auth_token=hf_api_token
-    )
-
-    logger.info("Serializing model")
-    logger.info("GPU: " + get_gpu_name())
-    logger.info("GPU RAM: " + get_vram_usage_str())
-    logger.info("PYTHON USED RAM: " + get_ram_usage_str())
-
-    serialize_model(
-        pipeline.text_encoder.eval(),
-        pipeline.text_encoder.config,
-        output_prefix,
-        "encoder",
-    )
-    serialize_model(pipeline.vae.eval(), None, output_prefix, "vae")
-    serialize_model(pipeline.unet.eval(), None, output_prefix, "unet")
-
-    pipeline.tokenizer.save_pretrained(output_prefix)
-
-    if args.validate:
-        device = get_device()
-
-        logger.info("Validating serialization")
-        vae = load_model(output_prefix, AutoencoderKL, None, "vae", device)
-        unet = load_model(output_prefix, UNet2DConditionModel, None, "unet")
-        encoder = load_model(
-            output_prefix, CLIPTextModel, CLIPTextConfig, "encoder", device
-        )
-
-        pipeline = StableDiffusionPipeline(
-            text_encoder=encoder,
-            vae=vae,
-            unet=unet,
-            tokenizer=CLIPTokenizer.from_pretrained(
-                args.input_directory, subfolder="tokenizer"
-            ),
-            scheduler=LMSDiscreteScheduler(
-                beta_end=0.012,
-                beta_schedule="scaled_linear",
-                beta_start=0.00085,
-                num_train_timesteps=1000,
-                trained_betas=None,
-            ),
-            safety_checker=None,
-            feature_extractor=None,
-        ).to(device)
-
-        prompt = "a photo of an astronaut riding a horse on mars"
-        with torch.autocast(
-            "cuda" if torch.cuda.is_available() else "cpu"
-        ):  # for some reason device_type needs to be a string instead of an actual device
-            pipeline(prompt).images[0]
-
-
-def hf_main(args):
-    output_prefix = args.output_prefix
-    print("MODEL PATH:", args.input_directory)
-    print("OUTPUT PREFIX:", output_prefix)
-
-    from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
-
-    model_config = AutoConfig.from_pretrained(args.input_directory)
-    model = AutoModelForCausalLM.from_pretrained(
-        args.input_directory, config=model_config, torch_dtype=torch.float16
-    )
-
-    logger.info("Serializing model")
-    logger.info("GPU: " + get_gpu_name())
-    logger.info("GPU RAM: " + get_vram_usage_str())
-    logger.info("PYTHON USED RAM: " + get_ram_usage_str())
-
-    serialize_model(model, model_config, output_prefix)
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.input_directory
-    ).save_pretrained(output_prefix)
-
-    if args.validate:
-        device = get_device()
-        logger.info("Validating serialization")
-        model = load_model(
-            output_prefix,
-            AutoModelForCausalLM,
-            AutoConfig,
-            None,
-            device,
-            "float16",
-        ).eval()
-        # test generation
-        tokenizer = AutoTokenizer.from_pretrained(args.input_directory)
-        input_ids = tokenizer.encode(
-            "¡Hola! Encantado de conocerte. hoy voy a", return_tensors="pt"
-        ).to(device)
-        with torch.no_grad():
-            output = model.generate(
-                input_ids, max_new_tokens=50, do_sample=True
-            )
-        logger.info(
-            f"Test Output: {tokenizer.decode(output[0], skip_special_tokens=True)}"
-        )
-
-
-def main():
-    # usage: tensorizer [input-directory] [output-prefix] [model-type]
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "input_directory",
-        type=str,
-        help="Path to model directory or HF model ID",
-    )
-    parser.add_argument(
-        "output_prefix", type=str, help="Path to output directory"
-    )
-    parser.add_argument(
-        "--model_type",
-        type=str,
-        choices=["transformers", "diffusers"],
-        required=True,
-        help="Framework used for the model",
-    )
-    parser.add_argument(
-        "--validate",
-        action="store_true",
-        help="Validate serialization by running a test inference",
-    )
-    args = parser.parse_args()
-
-    if args.model_type == "transformers":
-        hf_main(args)
-    elif args.model_type == "diffusers":
-        df_main(args)
-    else:
-        raise ValueError(
-            f"Unknown model type {args.model_type} (transformers or diffusers)"
-        )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tensorizer/tensors.proto b/tensorizer/tensors.proto
new file mode 120000
index 00000000..1aab8338
--- /dev/null
+++ b/tensorizer/tensors.proto
@@ -0,0 +1,1 @@
+../proto/tensors.proto
\ No newline at end of file
diff --git a/tensorizer/tensors_pb2.py b/tensorizer/tensors_pb2.py
new file mode 120000
index 00000000..e8100f6a
--- /dev/null
+++ b/tensorizer/tensors_pb2.py
@@ -0,0 +1,1 @@
+../tensors/tensors_pb2.py
\ No newline at end of file
diff --git a/tensorizer/utils.py b/tensorizer/utils.py
new file mode 100644
index 00000000..c1df5af3
--- /dev/null
+++ b/tensorizer/utils.py
@@ -0,0 +1,140 @@
+import torch
+
+try:
+    import resource
+except ImportError:
+    resource = None
+
+import psutil
+
+try:
+    import pynvml
+
+    try:
+        pynvml.nvmlInit()
+    except pynvml.nvml.NVMLError_LibraryNotFound:
+        pynvml = None
+except ImportError:
+    pynvml = None
+
+__all__ = [
+    "convert_bytes",
+    "get_device",
+    "get_mem_usage",
+    "get_gpu_name",
+    "no_init_or_tensor",
+]
+
+
+# Silly function to convert to human bytes
+def convert_bytes(num):
+    """
+    this function will convert bytes to MB.... GB... etc
+    """
+    step_unit = 1000.0
+
+    for x in ["bytes", "KB", "MB", "GB", "TB"]:
+        if num < step_unit:
+            return "%3.1f %s" % (num, x)
+        num /= step_unit
+
+
+def get_device() -> torch.device:
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _get_gpu_mem_usage():
+    """
+    Returns a tuple of (total_vram, free_vram, used_vram) in bytes if possible.
+    Returns None if neither PyTorch >=1.10 nor pynvml is available.
+    """
+    if not torch.cuda.is_available():
+        return None
+    mem_get_info = getattr(torch.cuda, "mem_get_info", None)
+    if mem_get_info is not None:
+        free, total = mem_get_info()
+        return total, free, total - free
+    elif pynvml is not None:
+        cudadev = torch.cuda.current_device()
+        nvml_device = pynvml.nvmlDeviceGetHandleByIndex(cudadev)
+        gpu_info = pynvml.nvmlDeviceGetMemoryInfo(nvml_device)
+        return gpu_info.total, gpu_info.free, gpu_info.used
+    else:
+        return None
+
+
+def get_mem_usage() -> str:
+    """
+    Returns memory usage statistics for the CPU, GPU, and Torch.
+
+    :return:
+    """
+    gpu_str = ""
+    torch_str = ""
+    try:
+        gpu_info = _get_gpu_mem_usage()
+        if gpu_info is not None:
+            gpu_total, gpu_free, gpu_used = gpu_info
+            gpu_total >>= 20
+            gpu_free >>= 20
+            gpu_used >>= 20
+            gpu_str = f"GPU: (U: {gpu_used:,}MiB F: {gpu_free:,}MiB T: {gpu_total:,}MiB) "
+        torch_reserved_gpu = torch.cuda.memory.memory_reserved() >> 20
+        torch_reserved_max = torch.cuda.memory.max_memory_reserved() >> 20
+        torch_used_gpu = torch.cuda.memory_allocated() >> 20
+        torch_max_used_gpu = torch.cuda.max_memory_allocated() >> 20
+        torch_str = (
+            f"TORCH: (R: {torch_reserved_gpu:,}MiB/"
+            f"{torch_reserved_max:,}MiB, "
+            f"A: {torch_used_gpu:,}MiB/{torch_max_used_gpu:,}MiB)"
+        )
+    except AssertionError:
+        pass
+    if resource is not None:
+        cpu_maxrss = (
+            resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+            + resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
+        ) >> 10
+    else:
+        process = psutil.Process()
+        cpu_maxrss = (
+            process.memory_info().rss
+            + sum(p.memory_info().rss for p in process.children(True))
+        ) >> 20
+    cpu_vmem = psutil.virtual_memory()
+    cpu_free = cpu_vmem.free >> 20
+    return (
+        f"CPU: (maxrss: {cpu_maxrss:,}MiB F: {cpu_free:,}MiB) "
+        f"{gpu_str}{torch_str}"
+    )
+
+
+def get_gpu_name() -> str:
+    if torch.cuda.is_available():
+        return torch.cuda.get_device_name(0)
+    return "N/A"
+
+
+def no_init_or_tensor(loading_code):
+    def dummy(self):
+        return
+
+    modules = [torch.nn.Linear, torch.nn.Embedding, torch.nn.LayerNorm]
+    original = {}
+    for mod in modules:
+        original[mod] = mod.reset_parameters
+        mod.reset_parameters = dummy
+    original_empty = torch.empty
+
+    # when torch.empty is called, make it map to meta device by replacing the
+    # device in kwargs.
+    torch.empty = lambda *args, **kwargs: original_empty(
+        *args, **{**kwargs, "device": "meta"}
+    )
+
+    result = loading_code()
+    for mod in modules:
+        mod.reset_parameters = original[mod]
+    torch.empty = original_empty
+
+    return result


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 00000000..9c70e049
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,23 @@
+name: Test
+
+on: [push, pull_request]
+
+jobs:
+  run_tests:
+    name: Run tests
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/checkout@v3
+
+      - uses: actions/setup-python@v3
+        with:
+          python-version: '3.8'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: >
+          python -m pip install -e . &&
+          python -m pip install -r tests/requirements.txt
+
+      - name: Run tests
+        run: python -m unittest discover tests/ --verbose
\ No newline at end of file
diff --git a/test_protobuf.py b/ATTIC/test_protobuf.py
similarity index 100%
rename from test_protobuf.py
rename to ATTIC/test_protobuf.py
diff --git a/test_tensorizer.py b/ATTIC/test_tensorizer.py
similarity index 100%
rename from test_tensorizer.py
rename to ATTIC/test_tensorizer.py
diff --git a/tests/requirements.txt b/tests/requirements.txt
new file mode 100644
index 00000000..5555bffc
--- /dev/null
+++ b/tests/requirements.txt
@@ -0,0 +1,2 @@
+transformers>=4.27.1
+moto[s3,server]>=4.1.4
\ No newline at end of file
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
new file mode 100644
index 00000000..788fe06a
--- /dev/null
+++ b/tests/test_serialization.py
@@ -0,0 +1,286 @@
+import contextlib
+import gc
+import os
+import tempfile
+import unittest
+import re
+from typing import Tuple
+import torch
+
+os.environ[
+    "TRANSFORMERS_VERBOSITY"
+] = "error"  # disable missing keys and unexpected key warnings
+os.environ[
+    "TOKENIZERS_PARALLELISM"
+] = "false"  # avoids excessive warnings about forking after using a tokenizer
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from tensorizer import TensorSerializer, TensorDeserializer
+from tensorizer import utils
+from collections import OrderedDict
+
+model_name = "EleutherAI/gpt-neo-125M"
+num_hellos = 400
+is_cuda_available = torch.cuda.is_available()
+default_device = "cuda" if is_cuda_available else "cpu"
+
+
+def serialize_model(model_name: str, device: str) -> Tuple[str, dict]:
+    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+    sd = model.state_dict()
+    out_file = tempfile.NamedTemporaryFile("wb+", delete=False)
+    try:
+        serializer = TensorSerializer(out_file)
+        serializer.write_module(model)
+        serializer.close()
+    except Exception:
+        os.unlink(out_file)
+        raise
+    return out_file.name, sd
+
+
+def check_deserialized(deserialized, model_name: str, allow_subset=False):
+    orig_sd = AutoModelForCausalLM.from_pretrained(model_name).state_dict()
+    if not allow_subset:
+        assert orig_sd.keys() == deserialized.keys()
+    for k, v in deserialized.items():
+        assert k in orig_sd
+        assert v.size() == orig_sd[k].size()
+        assert v.dtype == orig_sd[k].dtype
+        assert torch.all(orig_sd[k].to(v.device) == v)
+    del orig_sd
+    gc.collect()
+
+
+@contextlib.contextmanager
+def enable_tokenizers_parallelism():
+    os.environ["TOKENIZERS_PARALLELISM"] = "true"
+    try:
+        yield
+    finally:
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+def check_inference(
+    deserializer: TensorDeserializer, model_ref: str, device: str
+):
+    # This ensures that the model is not initialized.
+    model = utils.no_init_or_tensor(
+        lambda: AutoModelForCausalLM.from_pretrained(
+            model_ref, state_dict=OrderedDict()
+        )
+    )
+
+    deserializer.load_into_module(model)
+
+    # Tokenize and generate
+    with enable_tokenizers_parallelism():
+        tokenizer = AutoTokenizer.from_pretrained(model_ref)
+        input_ids = tokenizer.encode(
+            " hello" * num_hellos, return_tensors="pt"
+        ).to(device)
+
+        with torch.no_grad():
+            output = model.generate(
+                input_ids, max_new_tokens=50, do_sample=True
+            )
+
+        decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+        assert decoded.count("hello") > num_hellos
+
+
+class TestSerialization(unittest.TestCase):
+    def test_serialization(self):
+        for device in "cuda", "cpu":
+            if device == "cuda" and not is_cuda_available:
+                continue
+            with self.subTest(msg=f"Serializing with device {device}"):
+                gc.collect()
+                before_serialization = utils.get_mem_usage()
+                serialized_model, orig_sd = serialize_model(model_name, device)
+                after_serialization = utils.get_mem_usage()
+                print(f"Before serialization: {before_serialization}")
+                print(f"After serialization:  {after_serialization}")
+                del orig_sd
+                try:
+                    with open(serialized_model, "rb") as in_file:
+                        deserialized = TensorDeserializer(in_file, device="cpu")
+                        check_deserialized(deserialized, model_name)
+                        deserialized.close()
+                        del deserialized
+                finally:
+                    os.unlink(serialized_model)
+
+
+class TestDeserialization(unittest.TestCase):
+    _serialized_model_path: str
+
+    @classmethod
+    def setUpClass(cls):
+        serialized_model_path, sd = serialize_model(model_name, "cpu")
+        del sd
+        cls._serialized_model_path = serialized_model_path
+        gc.collect()
+
+    @classmethod
+    def tearDownClass(cls):
+        os.unlink(cls._serialized_model_path)
+
+    def test_default_cpu(self):
+        in_file = open(self._serialized_model_path, "rb")
+        gc.collect()
+        before_deserialization = utils.get_mem_usage()
+        deserialized = TensorDeserializer(in_file, device="cpu")
+        after_deserialization = utils.get_mem_usage()
+        check_deserialized(deserialized, model_name)
+        deserialized.close()
+        print(f"Before deserialization: {before_deserialization}")
+        print(f"After deserialization:  {after_deserialization}")
+
+    @unittest.skipIf(not is_cuda_available, "Requires CUDA")
+    def test_default_gpu(self):
+        in_file = open(self._serialized_model_path, "rb")
+        gc.collect()
+        before_deserialization = utils.get_mem_usage()
+        deserialized = TensorDeserializer(in_file, device="cuda")
+        check_deserialized(deserialized, model_name)
+        after_deserialization = utils.get_mem_usage()
+        deserialized.close()
+        print(f"Before deserialization: {before_deserialization}")
+        print(f"After deserialization:  {after_deserialization}")
+        del in_file
+        gc.collect()
+        after_del = utils.get_mem_usage()
+        print(f"After del: {after_del}")
+
+    def test_lazy_load(self):
+        in_file = open(self._serialized_model_path, "rb")
+        deserialized = TensorDeserializer(
+            in_file, device=default_device, lazy_load=True
+        )
+
+        check_deserialized(deserialized, model_name)
+        check_inference(deserialized, model_name, default_device)
+        deserialized.close()
+
+    @unittest.skipIf(not is_cuda_available, "plaid_mode requires CUDA")
+    def test_plaid_mode(self):
+        in_file = open(self._serialized_model_path, "rb")
+        deserialized = TensorDeserializer(
+            in_file, device="cuda", plaid_mode=True
+        )
+
+        check_deserialized(deserialized, model_name)
+        deserialized.close()
+
+    @unittest.skipIf(not is_cuda_available, "plaid_mode requires CUDA")
+    def test_plaid_mode_inference(self):
+        in_file = open(self._serialized_model_path, "rb")
+        deserialized = TensorDeserializer(
+            in_file, device="cuda", plaid_mode=True
+        )
+
+        check_inference(deserialized, model_name, "cuda")
+        deserialized.close()
+
+    @unittest.skipIf(not is_cuda_available, "plaid_mode requires CUDA")
+    def test_plaid_mode_guards(self):
+        in_file = open(self._serialized_model_path, "rb")
+        deserialized = TensorDeserializer(
+            in_file, device="cuda", plaid_mode=True
+        )
+        keys = list(deserialized.keys())
+        _ = deserialized[keys[0]]
+        _ = deserialized[keys[1]]
+
+        with self.assertRaises(RuntimeError):
+            _ = deserialized[keys[0]]
+
+        deserialized.close()
+
+    def test_s3(self):
+        deserialized = TensorDeserializer(
+            f"s3://tensorized/{model_name}/model.tensors", device=default_device
+        )
+        check_deserialized(deserialized, model_name)
+        check_inference(deserialized, model_name, default_device)
+        deserialized.close()
+
+    def test_s3_fp16(self):
+        deserialized = TensorDeserializer(
+            f"s3://tensorized/{model_name}/fp16/model.tensors",
+            device=default_device,
+        )
+        check_inference(deserialized, model_name, default_device)
+        deserialized.close()
+
+    def test_s3_lazy_load(self):
+        deserialized = TensorDeserializer(
+            f"s3://tensorized/{model_name}/model.tensors",
+            device=default_device,
+            lazy_load=True,
+        )
+        check_deserialized(deserialized, model_name)
+        check_inference(deserialized, model_name, default_device)
+        deserialized.close()
+
+    def test_filter_func(self):
+        # These two filters should produce identical results
+        pattern = re.compile(r"transformer\.h\.0.*")
+
+        def custom_check(tensor_name: str) -> bool:
+            return tensor_name.startswith("transformer.h.0")
+
+        # Testing no filter_func
+        in_file = open(self._serialized_model_path, "rb")
+        deserialized = TensorDeserializer(
+            in_file, device=default_device, filter_func=None
+        )
+        all_keys = set(deserialized.keys())
+        assert all_keys, (
+            "Deserializing the model with no filter_func"
+            " loaded an empty set of tensors"
+        )
+        check_deserialized(deserialized, model_name)
+        deserialized.close()
+
+        expected_regex_keys = set(filter(pattern.match, all_keys))
+        expected_custom_keys = set(filter(custom_check, all_keys))
+
+        assert (
+            expected_regex_keys
+            and expected_regex_keys < all_keys
+            and expected_custom_keys
+            and expected_custom_keys < all_keys
+        ), (
+            "The filter_func test cannot continue"
+            " because a filter_func used in the test"
+            " does not appear in the test model,"
+            " or matches all tensor names."
+            " Update the pattern and/or custom_check"
+            " to use more informative filtering criteria."
+            "\n\nTensors present in the model: " + " ".join(all_keys)
+        )
+
+        with self.subTest(msg="Testing regex filter_func"):
+            in_file = open(self._serialized_model_path, "rb")
+            deserialized = TensorDeserializer(
+                in_file, device=default_device, filter_func=pattern.match
+            )
+            regex_keys = set(deserialized.keys())
+            # Test that the deserialized tensors form a proper,
+            # non-empty subset of the original list of tensors.
+            assert regex_keys == expected_regex_keys
+            check_deserialized(deserialized, model_name, allow_subset=True)
+            deserialized.close()
+
+        with self.subTest(msg="Testing custom filter_func"):
+            in_file = open(self._serialized_model_path, "rb")
+            deserialized = TensorDeserializer(
+                in_file, device=default_device, filter_func=custom_check
+            )
+            custom_keys = set(deserialized.keys())
+            assert custom_keys == expected_custom_keys
+            check_deserialized(deserialized, model_name, allow_subset=True)
+            deserialized.close()
diff --git a/tests/test_stream_io.py b/tests/test_stream_io.py
new file mode 100644
index 00000000..16da70d9
--- /dev/null
+++ b/tests/test_stream_io.py
@@ -0,0 +1,186 @@
+import contextlib
+import os
+import unittest
+import moto
+import boto3
+from tensorizer import stream_io
+
+NEO_URL = (
+    "https://raw.githubusercontent.com/EleutherAI/gpt-neo/master/README.md"
+)
+
+
+class TestCurlStream(unittest.TestCase):
+    def test_curl_stream(self):
+        stream = stream_io.CURLStreamFile(NEO_URL)
+        self.assertEqual(b"# GPT Neo", stream.read(9))
+
+    def test_curl_stream_begin(self):
+        stream = stream_io.CURLStreamFile(NEO_URL, begin=2)
+        self.assertEqual(b"GPT Neo", stream.read(7))
+
+    def test_curl_stream_end(self):
+        stream = stream_io.CURLStreamFile(NEO_URL, end=2)
+        self.assertEqual(b"# ", stream.read(5))
+
+    def test_curl_stream_buffer_end(self):
+        stream = stream_io.CURLStreamFile(NEO_URL, end=2)
+        ba = bytearray(5)
+        self.assertEqual(2, stream.readinto(ba))
+        self.assertEqual(b"# \x00\x00\x00", ba)
+
+    def test_curl_stream_seek(self):
+        stream = stream_io.CURLStreamFile(NEO_URL)
+        stream.seek(2)
+        self.assertEqual(b"GPT Neo", stream.read(7))
+        stream.seek(6)
+        self.assertEqual(b"Neo\n", stream.read(4))
+
+    def test_curl_stream_seek_forward(self):
+        stream = stream_io.CURLStreamFile(NEO_URL)
+        stream.seek(2)
+        self.assertEqual(b"GPT Neo", stream.read(7))
+        stream.seek(13)
+        self.assertEqual(b"[DOI]", stream.read(5))
+
+
+def set_up_moto(*endpoints):
+    os.environ["MOTO_S3_CUSTOM_ENDPOINTS"] = ",".join(endpoints)
+
+    # Clear any environment variables that boto3 may attempt to access
+    # to avoid accidentally writing to a real bucket in case of failure
+    old_environment = {
+        key: os.environ[key]
+        for key in (
+            "AWS_ACCESS_KEY_ID",
+            "AWS_SECRET_ACCESS_KEY",
+            "AWS_SECURITY_TOKEN",
+            "AWS_SESSION_TOKEN",
+            "AWS_DEFAULT_REGION",
+        )
+        if key in os.environ
+    }
+
+    for key in old_environment:
+        test_value = "us-east-1" if key == "AWS_DEFAULT_REGION" else "TEST"
+        os.environ[key] = test_value
+
+    mock_s3 = moto.mock_s3()
+    access_key = mock_s3.FAKE_KEYS["AWS_ACCESS_KEY_ID"]
+    secret_key = mock_s3.FAKE_KEYS["AWS_SECRET_ACCESS_KEY"]
+    return mock_s3, access_key, secret_key, old_environment
+
+
+def tear_down_moto(old_environment):
+    for key, value in old_environment:
+        os.environ[key] = value
+
+
+@contextlib.contextmanager
+def mock_server():
+    from moto.server import ThreadedMotoServer
+
+    # Disable mock server logs
+    import logging
+
+    werkzeug_logger = logging.getLogger("werkzeug")
+    old_log_level = werkzeug_logger.getEffectiveLevel()
+    werkzeug_logger.setLevel(logging.CRITICAL + 1)
+
+    server = ThreadedMotoServer(
+        ip_address="127.0.0.1", port=5000, verbose=False
+    )
+    server.start()
+
+    # Disable https validation on endpoints
+    ensure_https_endpoint, stream_io._ensure_https_endpoint = (
+        stream_io._ensure_https_endpoint,
+        lambda endpoint: endpoint,
+    )
+    try:
+        yield "http://127.0.0.1:5000"
+    finally:
+        stream_io._ensure_https_endpoint = ensure_https_endpoint
+        server.stop()
+        werkzeug_logger.setLevel(old_log_level)
+
+
+class TestS3(unittest.TestCase):
+    endpoint: str
+    BUCKET_NAME: str
+    ACCESS_KEY: str
+    SECRET_KEY: str
+    old_environment: dict
+
+    @classmethod
+    def setUpClass(cls):
+        # Sets up a mock S3 environment with moto.
+        # Can be replaced for testing without mocks.
+        cls.endpoint = "https://" + stream_io.default_s3_write_endpoint
+        cls.BUCKET_NAME = "test-bucket"
+        (
+            cls.mock_s3,
+            cls.ACCESS_KEY,
+            cls.SECRET_KEY,
+            cls.old_environment,
+        ) = set_up_moto(cls.endpoint)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        tear_down_moto(cls.old_environment)
+
+    def setUp(self) -> None:
+        # This is in setUp/tearDown rather than setUpClass/tearDownClass
+        # so that mock bucket state is not shared between test runs
+        self.mock_s3.start()
+        s3 = boto3.resource("s3", endpoint_url=self.endpoint)
+        bucket = s3.Bucket(self.BUCKET_NAME)
+        bucket.create()
+
+    def tearDown(self) -> None:
+        self.mock_s3.stop()
+
+    def assert_bucket_contents(self, key, content):
+        # Not a test case
+        s3 = boto3.resource("s3")
+        obj = s3.Object(self.BUCKET_NAME, key)
+        actual = obj.get()["Body"].read()
+        self.assertEqual(actual, content)
+
+    def put_bucket_contents(self, key, content):
+        # Not a test case
+        s3 = boto3.resource("s3")
+        obj = s3.Object(self.BUCKET_NAME, key)
+        obj.put(Body=content)
+
+    def test_upload(self):
+        key = "model.tensors"
+        s = stream_io.open_stream(
+            f"s3://{self.BUCKET_NAME}/{key}",
+            mode="wb",
+            s3_access_key_id=self.ACCESS_KEY,
+            s3_secret_access_key=self.SECRET_KEY,
+            s3_endpoint=self.endpoint,
+        )
+        long_string = b"Hello" * 1024
+        s.write(long_string)
+        s.close()
+        self.assert_bucket_contents(key, long_string)
+
+    def test_download(self):
+        # CURLStreamFiles ignore moto's normal mocks, since moto still
+        # generates a real URL when a real endpoint is used, so instead
+        # we create an entire mock S3 server.
+        with mock_server() as endpoint:
+            key = "model.tensors"
+            long_string = b"Hello" * 1024
+            self.put_bucket_contents(key, long_string)
+            self.assert_bucket_contents(key, long_string)
+            with stream_io.open_stream(
+                f"s3://{self.BUCKET_NAME}/{key}",
+                mode="rb",
+                s3_access_key_id="X",
+                s3_secret_access_key="X",
+                s3_endpoint=endpoint,
+            ) as s:
+                self.assertEqual(s.read(), long_string)