[ghstack-poisoned]

[ghstack-poisoned]

[ghstack-poisoned]

[ghstack-poisoned]

[ghstack-poisoned]

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op device_mesh::_runtime_compute_coordinate_on_dim - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. DeviceMesh.sym_get_coordinate - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. Shard._select_split_tensor - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. Shard.local_shard_size_and_offset - Updated type hints to properly reflect SymInt return types.
5. Dynamo config - Always enables allow_dynamic_output_shape_ops (with TODO to reconsider if still needed).





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op device_mesh::_runtime_compute_coordinate_on_dim - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. DeviceMesh.sym_get_coordinate - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. Shard._select_split_tensor - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. Shard.local_shard_size_and_offset - Updated type hints to properly reflect SymInt return types.
5. Dynamo config - Always enables allow_dynamic_output_shape_ops (with TODO to reconsider if still needed).





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op device_mesh::_runtime_compute_coordinate_on_dim - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. DeviceMesh.sym_get_coordinate - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. Shard._select_split_tensor - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. Shard.local_shard_size_and_offset - Updated type hints to properly reflect SymInt return types.
5. Dynamo config - Always enables allow_dynamic_output_shape_ops (with TODO to reconsider if still needed).





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op device_mesh::_runtime_compute_coordinate_on_dim - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. DeviceMesh.sym_get_coordinate - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. Shard._select_split_tensor - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. Shard.local_shard_size_and_offset - Updated type hints to properly reflect SymInt return types.
5. Dynamo config - Always enables allow_dynamic_output_shape_ops (with TODO to reconsider if still needed).





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op device_mesh::_runtime_compute_coordinate_on_dim - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. DeviceMesh.sym_get_coordinate - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. Shard._select_split_tensor - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. Shard.local_shard_size_and_offset - Updated type hints to properly reflect SymInt return types.
5. Dynamo config - Always enables allow_dynamic_output_shape_ops (with TODO to reconsider if still needed).





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op device_mesh::_runtime_compute_coordinate_on_dim - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. DeviceMesh.sym_get_coordinate - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. Shard._select_split_tensor - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. Shard.local_shard_size_and_offset - Updated type hints to properly reflect SymInt return types.
5. Dynamo config - Always enables allow_dynamic_output_shape_ops (with TODO to reconsider if still needed).





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op `device_mesh::_runtime_compute_coordinate_on_dim` - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. `DeviceMesh.sym_get_coordinate` - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. `Shard._select_split_tensor` - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. `Shard.local_shard_size_and_offset` - Updated type hints to properly reflect SymInt return types.
5. Adds a config to enable compile_on_one_rank





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op `device_mesh::_runtime_compute_coordinate_on_dim` - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. `DeviceMesh.sym_get_coordinate` - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. `Shard._select_split_tensor` - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. `Shard.local_shard_size_and_offset` - Updated type hints to properly reflect SymInt return types.
5. Adds a config to enable compile_on_one_rank





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]

…ate computation in PT2"

This change enables compile-on-one-rank for DTensor slicing by making mesh coordinate lookups symbolic-aware.

1. New custom op `device_mesh::_runtime_compute_coordinate_on_dim` - An operator that computes mesh coordinates at runtime, allowing coordinate lookups to be deferred during tracing rather than baked in as constants.
2. `DeviceMesh.sym_get_coordinate` - Now uses the custom op when in fake mode (tracing), lifting the rank map as a graph constant and deferring the actual coordinate computation to runtime.
3. `Shard._select_split_tensor` - Extended to handle SymInt indices by using torch.narrow with symbolic start/length instead of list indexing, enabling symbolic tensor partitioning.
4. `Shard.local_shard_size_and_offset` - Updated type hints to properly reflect SymInt return types.
5. Adds a config to enable compile_on_one_rank





cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]