Attempt to make traceable allreduce work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
The long term solution should be "torchdispatch tensor subclass support", which is WIP

Option 1: Dynamo inlines Traceable Collectives, and handles AsyncTensor subclass
- AsyncTensor needs special handling since it is an unsupported subclass
- Special-case Dynamo to perform the simple function of AsyncTensor on its behalf:
  insert a 'Wait' op either "right after" the op that returned AsyncTensor (least optimal)
  or right before first user of AsyncTensor (optimal but more complicated)

  (a) add 'torch.distributed.traceable_collectives'
	  to `_dynamo.config.skipfiles_inline_module_allowlist`
  (b) implement AsyncTensor special behavior

Option 2: Dynamo allows traceable collectives in graph, doesn't inline
- we avoid Dynamo worrying about AsyncTensor at all. Just let it come up later,
  when doing dispatch tracing during AotAutograd
  (a) maybe the default already allows this op in graph!
  (b) still need AsyncTensor handling? Bc we seem to be faketensor tracing the op,
	  which fails when wrapping a FakeTensor in an AsyncTensor

[ghstack-poisoned]

Attempt to make traceable allreduce work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
The long term solution should be "torchdispatch tensor subclass support", which is WIP

Option 1: Dynamo inlines Traceable Collectives, and handles AsyncTensor subclass
- AsyncTensor needs special handling since it is an unsupported subclass
- Special-case Dynamo to perform the simple function of AsyncTensor on its behalf:
  insert a 'Wait' op either "right after" the op that returned AsyncTensor (least optimal)
  or right before first user of AsyncTensor (optimal but more complicated)

  (a) add 'torch.distributed.traceable_collectives'
	  to `_dynamo.config.skipfiles_inline_module_allowlist`
  (b) implement AsyncTensor special behavior

Option 2: Dynamo allows traceable collectives in graph, doesn't inline
- we avoid Dynamo worrying about AsyncTensor at all. Just let it come up later,
  when doing dispatch tracing during AotAutograd
  (a) maybe the default already allows this op in graph!
  (b) still need AsyncTensor handling? Bc we seem to be faketensor tracing the op,
	  which fails when wrapping a FakeTensor in an AsyncTensor

cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]

Attempt to make traceable allreduce work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
The long term solution should be "torchdispatch tensor subclass support", which is WIP

Option 1: Dynamo inlines Traceable Collectives, and handles AsyncTensor subclass
- AsyncTensor needs special handling since it is an unsupported subclass
- Special-case Dynamo to perform the simple function of AsyncTensor on its behalf:
  insert a 'Wait' op either "right after" the op that returned AsyncTensor (least optimal)
  or right before first user of AsyncTensor (optimal but more complicated)

  (a) add 'torch.distributed.traceable_collectives'
	  to `_dynamo.config.skipfiles_inline_module_allowlist`
  (b) implement AsyncTensor special behavior

Option 2: Dynamo allows traceable collectives in graph, doesn't inline
- we avoid Dynamo worrying about AsyncTensor at all. Just let it come up later,
  when doing dispatch tracing during AotAutograd
  (a) maybe the default already allows this op in graph!
  (b) still need AsyncTensor handling? Bc we seem to be faketensor tracing the op,
	  which fails when wrapping a FakeTensor in an AsyncTensor

cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]

Attempt to make traceable allreduce work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
The long term solution should be "torchdispatch tensor subclass support", which is WIP

Option 1: Dynamo inlines Traceable Collectives, and handles AsyncTensor subclass
- AsyncTensor needs special handling since it is an unsupported subclass
- Special-case Dynamo to perform the simple function of AsyncTensor on its behalf:
  insert a 'Wait' op either "right after" the op that returned AsyncTensor (least optimal)
  or right before first user of AsyncTensor (optimal but more complicated)

  (a) add 'torch.distributed.traceable_collectives'
	  to `_dynamo.config.skipfiles_inline_module_allowlist`
  (b) implement AsyncTensor special behavior

Option 2: Dynamo allows traceable collectives in graph, doesn't inline
- we avoid Dynamo worrying about AsyncTensor at all. Just let it come up later,
  when doing dispatch tracing during AotAutograd
  (a) maybe the default already allows this op in graph!
  (b) still need AsyncTensor handling? Bc we seem to be faketensor tracing the op,
	  which fails when wrapping a FakeTensor in an AsyncTensor

cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]

Attempt to make traceable allreduce work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
The long term solution should be "torchdispatch tensor subclass support", which is WIP

Option 1: Dynamo inlines Traceable Collectives, and handles AsyncTensor subclass
- AsyncTensor needs special handling since it is an unsupported subclass
- Special-case Dynamo to perform the simple function of AsyncTensor on its behalf:
  insert a 'Wait' op either "right after" the op that returned AsyncTensor (least optimal)
  or right before first user of AsyncTensor (optimal but more complicated)

  (a) add 'torch.distributed.traceable_collectives'
	  to `_dynamo.config.skipfiles_inline_module_allowlist`
  (b) implement AsyncTensor special behavior

Option 2: Dynamo allows traceable collectives in graph, doesn't inline
- we avoid Dynamo worrying about AsyncTensor at all. Just let it come up later,
  when doing dispatch tracing during AotAutograd
  (a) maybe the default already allows this op in graph!
  (b) still need AsyncTensor handling? Bc we seem to be faketensor tracing the op,
	  which fails when wrapping a FakeTensor in an AsyncTensor

cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]

Attempt to make traceable allreduce work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
The long term solution should be "torchdispatch tensor subclass support", which is WIP

Option 1: Dynamo inlines Traceable Collectives, and handles AsyncTensor subclass
- AsyncTensor needs special handling since it is an unsupported subclass
- Special-case Dynamo to perform the simple function of AsyncTensor on its behalf:
  insert a 'Wait' op either "right after" the op that returned AsyncTensor (least optimal)
  or right before first user of AsyncTensor (optimal but more complicated)

  (a) add 'torch.distributed.traceable_collectives'
	  to `_dynamo.config.skipfiles_inline_module_allowlist`
  (b) implement AsyncTensor special behavior

Option 2: Dynamo allows traceable collectives in graph, doesn't inline
- we avoid Dynamo worrying about AsyncTensor at all. Just let it come up later,
  when doing dispatch tracing during AotAutograd
  (a) maybe the default already allows this op in graph!
  (b) still need AsyncTensor handling? Bc we seem to be faketensor tracing the op,
	  which fails when wrapping a FakeTensor in an AsyncTensor

cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]

Make traceable collectives work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
For now, wait happens immediately, which generally forces an early sync.

Later, find a way either in dynamo or AOT stack to handle
AsyncCollectiveTensor to get the wait in the optimal place.

Note on implementation:
- Dynamo traces 'user-level' fc apis that are designed to behave differently
  in eager vs compiled.  In eager, there will be work-obj registration and
  a wrapper subclass will insert a 'wait' call at the appropriate time.
  In compile/trace mode, wait will be immetiately called, and work obj
  registration is required to be handled by the compile backend at runtime.
- Dynamo needs to trace into some of the helper functions in the 'user-level'
  api, such as '_expand_group' which is essentially a constant transformation.

[ghstack-poisoned]

Make traceable collectives work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
For now, wait happens immediately, which generally forces an early sync.

Later, find a way either in dynamo or AOT stack to handle
AsyncCollectiveTensor to get the wait in the optimal place.

Note on implementation:
- Dynamo traces 'user-level' fc apis that are designed to behave differently
  in eager vs compiled.  In eager, there will be work-obj registration and
  a wrapper subclass will insert a 'wait' call at the appropriate time.
  In compile/trace mode, wait will be immetiately called, and work obj
  registration is required to be handled by the compile backend at runtime.
- Dynamo needs to trace into some of the helper functions in the 'user-level'
  api, such as '_expand_group' which is essentially a constant transformation.

[ghstack-poisoned]

Make traceable collectives work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
For now, wait happens immediately, which generally forces an early sync.

Later, find a way either in dynamo or AOT stack to handle
AsyncCollectiveTensor to get the wait in the optimal place.

Note on implementation:
- Dynamo traces 'user-level' fc apis that are designed to behave differently
  in eager vs compiled.  In eager, there will be work-obj registration and
  a wrapper subclass will insert a 'wait' call at the appropriate time.
  In compile/trace mode, wait will be immetiately called, and work obj
  registration is required to be handled by the compile backend at runtime.
- Dynamo needs to trace into some of the helper functions in the 'user-level'
  api, such as '_expand_group' which is essentially a constant transformation.

[ghstack-poisoned]

Make traceable collectives work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
For now, wait happens immediately, which generally forces an early sync.

Later, find a way either in dynamo or AOT stack to handle
AsyncCollectiveTensor to get the wait in the optimal place.

Note on implementation:
- Dynamo traces 'user-level' fc apis that are designed to behave differently
  in eager vs compiled.  In eager, there will be work-obj registration and
  a wrapper subclass will insert a 'wait' call at the appropriate time.
  In compile/trace mode, wait will be immetiately called, and work obj
  registration is required to be handled by the compile backend at runtime.
- Dynamo needs to trace into some of the helper functions in the 'user-level'
  api, such as '_expand_group' which is essentially a constant transformation.

[ghstack-poisoned]

Make traceable collectives work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
For now, wait happens immediately, which generally forces an early sync.

Later, find a way either in dynamo or AOT stack to handle
AsyncCollectiveTensor to get the wait in the optimal place.

Note on implementation:
- Dynamo traces 'user-level' fc apis that are designed to behave differently
  in eager vs compiled.  In eager, there will be work-obj registration and
  a wrapper subclass will insert a 'wait' call at the appropriate time.
  In compile/trace mode, wait will be immetiately called, and work obj
  registration is required to be handled by the compile backend at runtime.
- Dynamo needs to trace into some of the helper functions in the 'user-level'
  api, such as '_expand_group' which is essentially a constant transformation.

[ghstack-poisoned]

Make traceable collectives work with torchdynamo,
bypassing problems with tracing the AsyncTensor subclass.

Accept a suboptimal solution for now, and optimize it later.
For now, wait happens immediately, which generally forces an early sync.

Later, find a way either in dynamo or AOT stack to handle
AsyncCollectiveTensor to get the wait in the optimal place.

Note on implementation:
- Dynamo traces 'user-level' fc apis that are designed to behave differently
  in eager vs compiled.  In eager, there will be work-obj registration and
  a wrapper subclass will insert a 'wait' call at the appropriate time.
  In compile/trace mode, wait will be immetiately called, and work obj
  registration is required to be handled by the compile backend at runtime.
- Dynamo needs to trace into some of the helper functions in the 'user-level'
  api, such as '_expand_group' which is essentially a constant transformation.

[ghstack-poisoned]