serializable data structure holding remote references to tensors

[casteryh]

Formulating some ideas from @allenwang28 .

It would be cool if we have some data structure satisfies the following

• Is serializable and can be sent through monarch rpc directly
• Small objects/non-tensor/metadata can be directly accessed upon deserialization
• Large objects/tensors are stored in torchstore and can be accessed via a materialize() / get() / fetch() method on demand.
• Easy to declare

Ideally we want to be able to do the following:

from torchstore.dataclasses import dataclass as ts_dataclass # a better name than `dataclass`?
from torchstore.dataclasses import RemoteTensor, mark_remote

@ts_dataclass
class Rollout:
    data: dict = {}
    some_int: int = 0
    # there is probably a way so that we don't have to explictly specify 'actions', might not worth the effort to do now
    actions: RemoteTensor | torch.tensor | torch.DTensor = mark_remote('actions')
    rewards: RemoteTensor | torch.tensor | torch.DTensor = mark_remote('rewards')
    envs: RemoteTensor | torch.tensor | torch.DTensor = mark_remote('envs')

Then we on the sending end we can do

rollout = Rollout(prefix='some_key_prefix') # or generate a random prefix for the user
# do something
rollout.data = some_trivially_serializable_python_object
rollout.actions = actions_tensor
# ...

# upload to torchstore and `dematerialize` the tensors before sending it out
rollout = await rollout.to_remote()
await other_actor.do_something.call(rollout)

On the receiving end we can do

# wait until everything is materialized
rollout = await rollout.materialize()

or

metadata = rollout.data
# do something with data
# pass to another actor without materializing
yet_another_actor.do_something_else.call(rollout)

or

actions = await rollout.data.materialize()
# do some computation with actions, envs is never fetched from remote

Other nice-to-haves:

• Support the above in a nested way, i.e., a field itself can have remote references to tensors.

[casteryh]

cc @joecummings @LucasLLC

[allenwang28]

Thanks @casteryh! Yeah this is a pretty good summary. A proposal for the nested interface is like:

@ts_dataclass
class ThisDataTypeWillBeRegistered:
    t1: torch.Tensor
    t2: torch.Tensor


@dataclass
class Container:
    name: str
    success: bool
    data: ThisDataTypeWillBeRegistered

then when Container is passed to different actors, the smaller data types like name, success are accessible, but I can instruct store to do stuff with the underlying data easily:

import torchstore as ts

class MyActor(Actor):
    @endpoint
    def test(self, c: Container):
        if c.success:
            c = await ts.get(c) # triggers the actual transfer over RDMA, presumably `data` to this point holds the key
        else:
            c = ts.del(c) # remove this entry

I would start with tensors as the basic "torch store MessageType" or something like that and build these custom types off of that

[allenwang28]

cc @pbontrager too

[LucasLLC]

My read from this is the following:

1. We want some support for complex object types
2. We want some concept of a storage 'future', which is a reference until accessed explicitly.

For now I think we should put off custom object types in store for as long as possible. That being said, I do think we need to treat dictionaries as first class citizens. What this means is if I do

ts.put(
    {
          'tensor': torch.Tensor,
          'obj': CustomObject
    },
    'state_dict'
)

I should be able to access via:

ts.get('state_dict'.tensor, tensor_slice=TensorSlice)
ts.get('state_dict'.obj)

We nee some amount of this for generic state dict support, but I think it's questionable if we should support nested state dict keys as the default case (rn we have a different api for pushing state dict).

On 2, I think it's reasonable to build these objects outside of torchstore if it significantly improves UX.