Enable eager spmd #7341
Enable eager spmd #7341
Conversation
JackCaoG
force-pushed
the
JackCaoG/eager_spmd
branch
from
3ec9f92 to
4244dcb
Compare
|
How could SPMD possibly work for eager mode? |
consider eager mode as calling |
Then how sharding propogation and auto partition work? I assume they don't carry states from last graph? |
The we will compile a graph for |
|
Okay, that's fair. |
| @@ -2742,7 +2742,9 @@ void XLANativeFunctions::_propagate_xla_data(const at::Tensor& input, | |||
|
|
|||
| // 2) Aid SPMD. | |||
| XLATensor::ShardingSpecPtr sharding = input_tensor->sharding_spec(); | |||
| if (sharding && sharding->sharding.type() != xla::OpSharding::UNKNOWN) { | |||
| // don't propagate sharding in eager mode. | |||
| if (!XLAGraphExecutor::Get()->UseEagerMode() && sharding && | |||
There was a problem hiding this comment.
It complained about the output tensor already has a sharding and we can't propagate to it. This happens in the backward. I didn't spend enough time to debug it but I don't expect user to actually run eager mode with step fn(forward and backward), I only expect them to run it with some data preprocessing on device so I just quickly unblock myself.
No description provided.