dmarx
commented
5 months ago - Absolutely! One strategy you can use is to partition the model into per-rank shards and serialize each shard into its own respective
.tensorsfile. Here's an example of what checkpointing with tensorizer could look like for distributed training, simplified from the coreweave/megatron fork
model = unwrap_model(model)
for i in range(len(model)):
mpu.set_virtual_pipeline_model_parallel_rank(i)
shard_serializer = TensorSerializer(
f"{checkpoint_name}_shard{i}.tensors"
)
shard_serializer.write_module(model[i])
shard_serializer.close()And the corresponding deserialization, again simplified from the linked code:
for i in range(len(model)):
mpu.set_virtual_pipeline_model_parallel_rank(i)
shard_deserializer = TensorDeserializer(
f"{checkpoint_name}_shard{i}.tensors",
device=torch.cuda.current_device(),
plaid_mode=True
)
shard_deserializer.load_into_module(model[i])
shard_deserializer.close()Alternatively, if you have a checkpoint already that you want to load across multiple ranks, you can provide a filter function to indicate which tensors should or should not be streamed into any given rank. EDIT: I just learned this is actually a pretty inefficient approach. If you have sufficient RAM, the lowest latency approach will be to deserialize the entire checkpoint onto your CPU, then assign tensors to devices as you normally would from there.
# EDIT: Actually, don't do this. leaving it here for demonstrative purposes.
for i in range(len(model)):
module_names = modules_to_load_into_rank[i]
# Should evaluate to True when the provided `key`
# matches the name of a tensor we want on the device
def filter_func(key: str) -> Callable[[str], bool]:
return key in module_names
mpu.set_virtual_pipeline_model_parallel_rank(i)
shard_deserializer = TensorDeserializer(
f"{checkpoint_name}_shard{i}.tensors",
device=torch.cuda.current_device(),
plaid_mode=True
)
shard_deserializer.load_into_module(model[i], filter_func)
shard_deserializer.close().tensorsfiles pack tensors into contiguous memory on disk. The chunk for a given tensor starts with a header containing metadata, and then is followed by the tensor's data. Main implementation details can be found here The data format doesn't do sharding or distribution on its own, but as demonstrated above is compatible with user-specified sharding. Alternatively, Coreweave storage supports both CephFS and Vast distributed file systems. We are currently undergoing benchmarking exercises whose results we will hopefully share soon. Based on our preliminary observations, reading and writing shards distributed across multiple.tensorsfiles is likely to be faster than one big file when using one of those distributed file systems.
Did that answer your questions?
SujoyDutta
I'm currently evaluating Tensorizer for handling large models, specifically models with parameters as larger than 70B that cannot be fit into a single GPU.
I have a few questions and concerns regarding Tensorizer's support for such large models and its handling of
.tensorsfiles.Large Model Support: Can Tensorizer effectively handle large models with parameters in the range of 70B and above wiht multiple distributed GPUs ? I havent seen any benchmarks for streaming large models that cannot be fit into a
single A100 GPU. Specifically, can Tensorizer scale effectively to accommodate large models across distributed environments, such as distributed training setups with multiple GPUs or nodes ?.tensors File Format: Could you provide some insights into the .tensors file format used by Tensorizer? How does Tensorizer organize and store tensors within .tensors files, especially for large models? Does it employ any sharding or distribution mechanisms for managing large tensors?
I appreciate any insights or documentation you can provide regarding these questions. Understanding Tensorizer's capabilities and limitations will help us make informed decisions for our use case.