torch.distributed.rpc package not work well with generator and lambda

[frank-dong-ms-zz]

I'm using torch.distributed.rpc package to work on a distributed training POC, currently I'm seeing rpc package itself is using pickle and pickle not work well with some python features like generator and lambda, so that put extra limitation, what even worse is if I try to reference some other python package that used lambda or generator I have no way to combine usage of this and rpc package.

So my question is:

1. can we use other pickling package that has fewer limitations like dill?
2. Any suggestion to work around this issue, especially when I try to use rpc with another package already using lambda or generator?

Thanks for any suggestions.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528 @jjlilley @osalpekar @jiayisuse @agolynski

[mrshenli]

Hey @frank-dong-ms

can we use other pickling package that has fewer limitations like dill?

It might be hard to pull dill in as a dependency for PyTorch, but applications should be able to override the pickler used by RPC. See the code below:

pytorch/torch/distributed/rpc/internal.py

Lines 140 to 149 in 7c07c39

	# Create _internal_rpc_pickler only once to initialize _dispatch_table only once
	_internal_rpc_pickler = _InternalRPCPickler()
	def serialize(obj):
	return _internal_rpc_pickler.serialize(obj)
	def deserialize(binary_data, tensor_table):
	return _internal_rpc_pickler.deserialize(binary_data, tensor_table)

Any suggestion to work around this issue, especially when I try to use rpc with another package already using lambda or generator?

Could you share an example that can repro the serde error?