问题类型

功能请求

来源

二进制文件

Tensorflow版本

2.10

自定义代码

无

OS平台和发行版

• 无响应*

移动设备

• 无响应*

Python版本

• 无响应*

Bazel版本

• 无响应*

GCC/编译器版本

• 无响应*

CUDA/cuDNN版本

• 无响应*

GPU型号和内存

• 无响应*

当前行为？

CollectiveOpV3Kernel did not support group assignments, I need a all2all OP which could dispatch tensors with different shape like Horovod.

重现问题的独立代码

import os

import tensorflow as tf
from tensorflow.python.ops import collective_ops

tf.config.experimental.set_memory_growth=True

group_size = 2
group_key = 104
device = 'GPU'
communication = 'auto'

dev0 = '/device:%s:0' % device
dev1 = '/device:%s:1' % device

@tf.function
def run_all_to_all_2devices():
  collectives = []
  in0 = tf.convert_to_tensor([1,1,1,1,3,3])
  in1 = tf.convert_to_tensor([2,2,2,4])
  with tf.device(dev0):
    group_handle0 = collective_ops.initialize_communicator(
        group_key=group_key,
        rank=0,
        group_size=group_size,
        communication_hint=communication)
    collectives.append(
        collective_ops.all_to_all_v3(group_handle0, in0, [0,4]))
  with tf.device(dev1):
    group_handle1 = collective_ops.initialize_communicator(
        group_key=group_key,
        rank=1,
        group_size=group_size,
        communication_hint=communication)
    collectives.append(
        collective_ops.all_to_all_v3(group_handle1, in1, [0,3]))
  return collectives

result = run_all_to_all_2devices()
print(result[0])
print(result[1])

相关日志输出

UnimplementedError                        Traceback (most recent call last)
Cell In[6], line 32
     28     collectives.append(
     29         collective_ops.all_to_all_v3(group_handle1, in1, [0,3]))
     30   return collectives
---> 32 result = run_all_to_all_2devices()
     33 print(result[0])
     34 print(result[1])

File ~/.conda/envs/hj/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File ~/.conda/envs/hj/lib/python3.8/site-packages/tensorflow/python/eager/execute.py:54, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52 try:
     53   ctx.ensure_initialized()
---> 54   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                       inputs, attrs, num_outputs)
     56 except core._NotOkStatusException as e:
     57   if name is not None:
...
	 [[CollectiveAllToAllV3_1/_11]]
  (1) UNIMPLEMENTED:  Group assignments are not supported yet.
	 [[{{node CollectiveAllToAllV3_1}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_run_all_to_all_2devices_39]