Repository: DeepRec-AI/HybridBackend (language: python)

Issue:
merge embedding table 
<img width="1634" alt="image" src="https://user-images.githubusercontent.com/33950866/216267372-cf1ceae2-4499-4511-bddc-1660d6ee4705.png">

Too many alltoall op to performance degradation


Discussion / hints:

Fix patch:
diff --git a/Makefile b/Makefile
index ef68aa4d..b6deedfd 100644
--- a/Makefile
+++ b/Makefile
@@ -23,7 +23,6 @@ HYBRIDBACKEND_DEBUG ?= OFF
 HYBRIDBACKEND_WHEEL_ALIAS ?=
 HYBRIDBACKEND_WHEEL_BUILD ?=
 HYBRIDBACKEND_WHEEL_REQUIRES ?=
-HYBRIDBACKEND_WHEEL_DEBUG ?= OFF
 HYBRIDBACKEND_WHEEL_REPAIR ?= ON
 HYBRIDBACKEND_WHEEL_POSTCHECK ?= ON
 HYBRIDBACKEND_CHECK_INSTANCE ?= OFF
@@ -239,7 +238,6 @@ build: $(CORE_DEPS)
 	WHEEL_ALIAS="$(HYBRIDBACKEND_WHEEL_ALIAS)" \
 	WHEEL_BUILD="$(HYBRIDBACKEND_WHEEL_BUILD)" \
 	WHEEL_REQUIRES="$(HYBRIDBACKEND_WHEEL_REQUIRES)" \
-	WHEEL_DEBUG="$(HYBRIDBACKEND_WHEEL_DEBUG)" \
 	$(PYTHON) setup.py bdist_wheel -d build/wheel
 	@ls build/wheel/*.whl
 
diff --git a/docs/tutorial/ranking/criteo/train.py b/docs/tutorial/ranking/criteo/train.py
index 795bc11e..72db6418 100755
--- a/docs/tutorial/ranking/criteo/train.py
+++ b/docs/tutorial/ranking/criteo/train.py
@@ -162,7 +162,8 @@ def main(args):
   if args.eval_every_n_iter is not None:
     hooks.append(hb.train.EvaluationHook(
       lambda: model.evaluate(eval_filenames),
-      every_n_iter=args.eval_every_n_iter))
+      every_n_iter=args.eval_every_n_iter,
+      summary_dir=args.output_dir))
   if args.log_every_n_iter is not None:
     hooks.append(
       tf.train.LoggingTensorHook(
diff --git a/docs/tutorial/ranking/taobao/train.py b/docs/tutorial/ranking/taobao/train.py
index 52eea918..65c4e428 100755
--- a/docs/tutorial/ranking/taobao/train.py
+++ b/docs/tutorial/ranking/taobao/train.py
@@ -144,7 +144,8 @@ def main(args):
   if args.eval_every_n_iter is not None:
     hooks.append(hb.train.EvaluationHook(
       lambda: model.evaluate(eval_filenames),
-      every_n_iter=args.eval_every_n_iter))
+      every_n_iter=args.eval_every_n_iter,
+      summary_dir=args.output_dir))
   if args.log_every_n_iter is not None:
     hooks.append(
       tf.train.LoggingTensorHook(
diff --git a/hybridbackend/tensorflow/__init__.py b/hybridbackend/tensorflow/__init__.py
index e3b22def..1cae72c4 100644
--- a/hybridbackend/tensorflow/__init__.py
+++ b/hybridbackend/tensorflow/__init__.py
@@ -20,9 +20,11 @@
 from __future__ import division
 from __future__ import print_function
 
+import contextlib as _ctxlib
+
 from hybridbackend.libhybridbackend import buildinfo
-from hybridbackend.tensorflow.feature_column.dense_features import \
-  dense_features
+from hybridbackend.tensorflow.framework.config import get_session_config
+from hybridbackend.tensorflow.framework.config import wraps_session_config
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.context import context
 from hybridbackend.tensorflow.framework.rewriting import function
@@ -33,10 +35,18 @@
 from . import data
 from . import distribute
 from . import estimator
-from . import feature_column
 from . import keras
 from . import metrics
 from . import plugins
 from . import training as train
 
 __version__ = buildinfo()
+
+
+@_ctxlib.contextmanager
+def embedding_scope(**kwargs):
+  r'''Scope for defining embedding weights.
+  '''
+  kwargs.setdefault('sharding', True)
+  with scope(**kwargs) as ctx:
+    yield ctx
diff --git a/hybridbackend/tensorflow/benchmarks/collective_benchmark.py b/hybridbackend/tensorflow/benchmarks/collective_benchmark.py
new file mode 100644
index 00000000..a303493f
--- /dev/null
+++ b/hybridbackend/tensorflow/benchmarks/collective_benchmark.py
@@ -0,0 +1,205 @@
+# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+r'''Collective benchmark.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import time
+
+import numpy as np
+import tensorflow as tf
+
+import hybridbackend.tensorflow as hb
+
+
+# pylint: disable=missing-docstring
+def allreduce(message_floats, message_partitions, message_device, topology):
+  del topology
+  step = tf.train.get_or_create_global_step()
+  results = [step.assign_add(1)]
+  for p in range(message_partitions):
+    with tf.device(f'/{message_device}:0'):
+      coll_input = tf.get_variable(
+        f'input{p}',
+        initializer=tf.random.normal(
+          [int(message_floats / message_partitions)],
+          mean=100,
+          stddev=80))
+    coll_output = hb.distribute.allreduce(coll_input)
+    with tf.device(f'/{message_device}:0'):
+      results.append(tf.identity(coll_output))
+  return tf.group(results)
+
+
+def alltoall(message_floats, message_partitions, message_device):
+  del topology
+  step = tf.train.get_or_create_global_step()
+  results = [step.assign_add(1)]
+  for p in range(message_partitions):
+    with tf.device(f'/{message_device}:0'):
+      coll_input = tf.get_variable(
+        f'input{p}',
+        initializer=tf.random.normal(
+          [int(message_floats / message_partitions)],
+          mean=100,
+          stddev=80))
+    coll_output = hb.distribute.alltoall(coll_input)
+    with tf.device(f'/{message_device}:0'):
+      results.append(tf.identity(coll_output))
+  return tf.group(results)
+
+
+def _uniform_sizes(total_size, active_size):
+  uniform_distro = [
+    1. / active_size for _ in range(active_size)]
+  return np.random.multinomial(total_size, uniform_distro, size=1)[0].tolist()
+
+
+def alltoallv(
+    message_floats, message_partitions, message_device, topology,
+    random_sizes=True):
+  active_size = hb.distribute.active_size(topology)
+  step = tf.train.get_or_create_global_step()
+  results = [step.assign_add(1)]
+  for p in range(message_partitions):
+    with tf.device(f'/{message_device}:0'):
+      message_count = int(message_floats / message_partitions)
+      coll_input = tf.get_variable(
+        f'input{p}',
+        initializer=tf.random.normal(
+          [message_count],
+          mean=100,
+          stddev=80))
+      if random_sizes:
+        coll_input_sizes = tf.constant(
+          _uniform_sizes(message_count, active_size),
+          dtype=tf.int32)
+      else:
+        message_divided = message_count // active_size
+        coll_input_sizes = tf.constant(
+          [message_divided for _ in range(active_size)],
+          dtype=tf.int32)
+    coll_output, _ = hb.distribute.alltoall(
+      coll_input, sizes=coll_input_sizes, topology=topology)
+    with tf.device(f'/{message_device}:0'):
+      results.append(tf.identity(coll_output))
+  return tf.group(results)
+
+
+def benchmark(args):
+  collective_ops = {
+    'allreduce': allreduce,
+    'alltoall': alltoall,
+    'alltoallv_': lambda mf, mp, md, topo: alltoallv(
+      mf, mp, md, topo, random_sizes=False),
+    'alltoallv': alltoallv}
+  for cop in args.collective_ops:
+    if cop not in collective_ops:
+      raise ValueError(
+        f'Specified collective op type `{cop}` '
+        f'not in {collective_ops.keys()}')
+  for nf in args.message_floats:
+    if nf % hb.context.world_size != 0:
+      raise ValueError(
+        f'#floats {nf} cannot be divided onto {hb.context.world_size} devices')
+    for part in args.message_partitions:
+      if nf % part != 0:
+        raise ValueError(
+          f'#floats {nf} cannot be divided into {part} partitions')
+
+  with tf.Graph().as_default(), hb.scope():
+    bench_ops = {
+      cop: {
+        nf: {
+          part: {
+            dev: {} for dev in args.message_devices}
+          for part in args.message_partitions}
+        for nf in args.message_floats}
+      for cop in args.collective_ops}
+    for cop in args.collective_ops:
+      with tf.name_scope(cop), tf.variable_scope(cop):
+        for nf in args.message_floats:
+          with tf.name_scope(f'{nf}floats'), tf.variable_scope(f'{nf}floats'):
+            for p in args.message_partitions:
+              with tf.name_scope(f'{p}parts'), tf.variable_scope(f'{p}parts'):
+                for dev in args.message_devices:
+                  with tf.name_scope(f'{dev}'), tf.variable_scope(f'{dev}'):
+                    for topo in args.collective_topology:
+                      with tf.name_scope(f'{topo}topology'), tf.variable_scope(
+                          f'{topo}topology'):
+                        bench_ops[cop][nf][p][dev][topo] = (
+                          collective_ops[cop](
+                            nf, p, dev, topo))
+    with tf.train.MonitoredTrainingSession('') as sess:
+      print('Rank\tCollective\tTopology\tDevice\tSize\t#Splits\tThroughput')
+      # pylint: disable=too-many-nested-blocks
+      for cop in args.collective_ops:
+        for dev in args.message_devices:
+          for nf in args.message_floats:
+            for part in args.message_partitions:
+              for topo in args.collective_topology:
+                for _ in range(args.warmup_steps):
+                  sess.run(bench_ops[cop][nf][part][dev][topo])
+                prev_ts = time.time()
+                for _ in range(args.num_steps):
+                  sess.run(bench_ops[cop][nf][part][dev][topo])
+                duration = time.time() - prev_ts
+                message_mbs = nf * 4. / 1024. / 1024.
+                print(
+                  f'{hb.context.rank}/{hb.context.world_size}\t'
+                  f'{cop}\tTopology-{topo}\t{dev}\t'
+                  f'{message_mbs:.2f}MB\t{part}\t'
+                  f'{args.num_steps * message_mbs * 8.0 / duration:.2f}Gb/s')
+
+
+if __name__ == '__main__':
+  tf.logging.set_verbosity(tf.logging.INFO)
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+    '--collective-ops',
+    nargs='+',
+    help='Collective ops in (allreduce, alltoall, alltoallv)',
+    default=['allreduce', 'alltoall', 'alltoallv'])
+  parser.add_argument(
+    '--collective-topology',
+    type=int,
+    nargs='+',
+    help='All/Intra/Inter nodes participate collective_ops',
+    default=[0, 1, 2])
+  parser.add_argument(
+    '--message-floats',
+    type=int,
+    nargs='+',
+    help='Count of floats in each message',
+    default=[65536, 262144, 1048576, 4194304, 16777216])
+  parser.add_argument(
+    '--message-partitions',
+    type=int,
+    nargs='+',
+    help='Number of partitions of each message',
+    default=[1, 8, 64])
+  parser.add_argument(
+    '--message-devices',
+    nargs='+',
+    help='Number of devices of each message',
+    default=['gpu', 'cpu'])
+  parser.add_argument('--warmup-steps', type=int, default=100)
+  parser.add_argument('--num-steps', type=int, default=500)
+  benchmark(parser.parse_args())
diff --git a/hybridbackend/tensorflow/benchmarks/embedding_benchmark_tier1.py b/hybridbackend/tensorflow/benchmarks/embedding_benchmark_tier1.py
new file mode 100644
index 00000000..5947b498
--- /dev/null
+++ b/hybridbackend/tensorflow/benchmarks/embedding_benchmark_tier1.py
@@ -0,0 +1,100 @@
+# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+r'''Embedding training using unobtrusive API on single GPU benchmark.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import os
+import time
+
+import tensorflow as tf
+
+import hybridbackend.tensorflow as hb
+
+
+# pylint: disable=missing-docstring
+def benchmark(params):
+  with tf.Graph().as_default():
+    fields = hb.data.Dataset.schema_from_parquet(params.filenames[0])
+    fields = [
+      f for f in fields
+      if f.name not in ('label', 'ts') and f.dtype in (tf.int32, tf.int64)]
+    ds = hb.data.Dataset.from_parquet(params.filenames, fields=fields)
+    ds = ds.batch(params.batch_size, drop_remainder=True)
+    ds = ds.prefetch(1)
+    iterator = tf.data.make_one_shot_iterator(ds)
+    iterator = hb.data.Iterator(iterator)
+    iterator_hook = hb.data.Iterator.Hook()
+    inputs = iterator.get_next()
+    outputs = []
+    with tf.name_scope('features'):
+      for field in inputs:
+        with tf.device('/cpu:0'):
+          embedding_weights = tf.get_variable(
+            f'{field}_weight',
+            shape=(128, 32),
+            initializer=tf.random_uniform_initializer(-1e-3, 1e-3))
+        ids = inputs[field]
+        if isinstance(ids, tf.Tensor):
+          ids = ids % params.dimension_size
+          ids, idx = tf.unique(tf.reshape(ids, shape=[-1]))
+          embeddings = tf.nn.embedding_lookup(embedding_weights, ids)
+          embeddings = tf.gather(embeddings, idx)
+        else:
+          ids = tf.SparseTensor(
+            ids.indices,
+            ids.values % params.dimension_size,
+            ids.dense_shape)
+          embeddings = tf.nn.embedding_lookup_sparse(
+            embedding_weights, ids, None)
+        outputs.append(embeddings)
+    loss = tf.math.add_n([tf.reduce_sum(t) for t in outputs])
+    opt = hb.train.AdagradOptimizer(learning_rate=0.01)
+    step = tf.train.get_or_create_global_step()
+    train_op = opt.minimize(loss, global_step=step)
+
+    with tf.train.MonitoredTrainingSession('', hooks=[iterator_hook]) as sess:
+      count = 0
+      prev_ts = time.time()
+      try:
+        while not sess.should_stop():
+          sess.run(train_op)
+          count += 1
+      except tf.errors.OutOfRangeError:
+        pass
+      duration = time.time() - prev_ts
+      if count <= 0:
+        print('Training embedding layers stopped unexpectedly')
+        return
+      print(
+        'Training embedding layers elapsed in '
+        f'{params.batch_size * count / duration:.2f} samples/sec ('
+        f'{1000. * duration / count:.2f} msec/step)')
+
+
+if __name__ == '__main__':
+  os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+  os.environ['HB_OP_RELOCATION_ENABLED'] = '1'
+  tf.logging.set_verbosity(tf.logging.INFO)
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--batch-size', type=int, default=16384)
+  parser.add_argument('--dimension-size', type=int, default=32)
+  parser.add_argument('filenames', nargs='+')
+  benchmark(parser.parse_args())
diff --git a/hybridbackend/tensorflow/benchmarks/partition_by_dual_modulo_benchmark.py b/hybridbackend/tensorflow/benchmarks/partition_by_dual_modulo_benchmark.py
new file mode 100644
index 00000000..9bf0df4d
--- /dev/null
+++ b/hybridbackend/tensorflow/benchmarks/partition_by_dual_modulo_benchmark.py
@@ -0,0 +1,123 @@
+# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+r'''Benchmark for partition, lookup and stitch.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import os
+from six.moves import xrange  # pylint: disable=redefined-builtin
+
+import tensorflow as tf
+
+import hybridbackend.tensorflow as hb
+
+
+# pylint: disable=missing-docstring
+def shard_lookup(weight_list, ids_list, params, fusion):
+  results = []
+  if fusion == 'disabled':
+    for i, weight in enumerate(weight_list):
+      with tf.name_scope(f'input_{i}'):
+        ids = ids_list[i]
+        ids_shards, idx_shards = hb.distribute.partition(
+          ids, params.num_partitions)
+        embs_shards = [
+          tf.gather(weight, ids_shards[i] % params.bucket_size)
+          for i in xrange(params.num_partitions)]
+        results.append(tf.dynamic_stitch(idx_shards, embs_shards))
+  elif fusion == 'vertical_and_horizontal_stage_one':
+    for i, weight in enumerate(weight_list):
+      with tf.name_scope(f'input_{i}'):
+        ids = ids_list[i]
+        ids, _, idx = hb.distribute.partition_by_dual_modulo_stage_one(
+          ids, params.num_partitions, params.modulus)
+        embs = tf.gather(weight, ids % params.bucket_size)
+        results.append(tf.gather(embs, idx))
+  else:
+    for i, weight in enumerate(weight_list):
+      with tf.name_scope(f'input_{i}'):
+        ids = ids_list[i]
+        ids, _, idx = hb.distribute.partition_by_dual_modulo_stage_two(
+          ids, params.num_partitions, params.modulus)
+        embs = tf.gather(weight, ids % params.bucket_size)
+        results.append(tf.gather(embs, idx))
+  return results
+
+
+def build_bench_op(params):
+  value_limit = 3 * params.num_columns * params.column_ids
+  with tf.device(params.device):
+    weight_initializer = tf.random_uniform(
+      [params.bucket_size, params.dim_size], 0, 100.0,
+      dtype=tf.float32)
+    weights = [
+      tf.get_variable(
+        f'weight{c}',
+        use_resource=False,
+        dtype=tf.float32,
+        initializer=weight_initializer)
+      for c in xrange(params.num_columns)]
+  input_initializer = tf.random_uniform(
+    [params.column_ids], 0, value_limit, dtype=tf.int32)
+  inputs = [
+    tf.get_variable(
+      f'input{c}',
+      use_resource=False,
+      dtype=tf.int32,
+      initializer=input_initializer)
+    for c in xrange(params.num_columns)]
+  bench_ops = shard_lookup(
+    weights, inputs, params,
+    fusion=params.fusion)
+  bench_ops = [tf.math.reduce_sum(o) for o in bench_ops]
+  step = tf.train.get_or_create_global_step()
+  bench_ops.append(tf.cast(step.assign_add(1), tf.float32))
+  bench_op = tf.math.add_n(bench_ops)
+  return bench_op
+
+
+@hb.function()
+def benchmark(params):
+  bench_op = build_bench_op(params)
+  hooks = [tf.train.StopAtStepHook(params.num_steps)]
+  hooks.append(
+    hb.train.StepStatHook(
+      count=params.column_ids * params.num_columns / 1000000.,
+      unit='Mlookups'))
+  with tf.train.MonitoredTrainingSession('', hooks=hooks) as sess:
+    while not sess.should_stop():
+      sess.run(bench_op)
+
+
+if __name__ == '__main__':
+  os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+  tf.logging.set_verbosity(tf.logging.INFO)
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--device', type=str, default='/gpu:0')
+  parser.add_argument('--column-ids', type=int, default=100000)
+  parser.add_argument('--dim-size', type=int, default=32)
+  parser.add_argument('--bucket-size', type=int, default=10000)
+  parser.add_argument('--num-columns', type=int, default=100)
+  parser.add_argument('--num-partitions', type=int, default=8)
+  parser.add_argument('--modulus', type=int, default=2)
+  parser.add_argument('--num-steps', type=int, default=10)
+  parser.add_argument(
+    '--fusion', type=str, default='vertical_and_horizontal_stage_two')
+  benchmark(parser.parse_args())
diff --git a/hybridbackend/tensorflow/ops/floormod_shuffle/benchmark.py b/hybridbackend/tensorflow/benchmarks/partition_by_modulo_benchmark.py
similarity index 89%
rename from hybridbackend/tensorflow/ops/floormod_shuffle/benchmark.py
rename to hybridbackend/tensorflow/benchmarks/partition_by_modulo_benchmark.py
index a5708b56..df3f8068 100644
--- a/hybridbackend/tensorflow/ops/floormod_shuffle/benchmark.py
+++ b/hybridbackend/tensorflow/benchmarks/partition_by_modulo_benchmark.py
@@ -36,27 +36,20 @@ def shard_lookup(weight_list, ids_list, params, fusion):
     for i, weight in enumerate(weight_list):
       with tf.name_scope(f'input_{i}'):
         ids = ids_list[i]
-        ids_shards, idx_shards = hb.floormod_partition(
+        ids_shards, idx_shards = hb.distribute.partition(
           ids, params.num_partitions)
         embs_shards = [
           tf.gather(weight, ids_shards[i] % params.bucket_size)
           for i in xrange(params.num_partitions)]
         results.append(tf.dynamic_stitch(idx_shards, embs_shards))
-  elif fusion == 'vertical':
+  else:
     for i, weight in enumerate(weight_list):
       with tf.name_scope(f'input_{i}'):
         ids = ids_list[i]
-        ids, _, idx = hb.floormod_shuffle(ids, params.num_partitions)
+        ids, _, idx = hb.distribute.partition_by_modulo(
+          ids, params.num_partitions)
         embs = tf.gather(weight, ids % params.bucket_size)
         results.append(tf.gather(embs, idx))
-  else:
-    ids_list, _, idx_list = hb.floormod_shuffle_n(
-      ids_list, params.num_partitions)
-    results = []
-    for i, weight in enumerate(weight_list):
-      with tf.name_scope(f'input_{i}'):
-        embs = tf.gather(weight, ids_list[i] % params.bucket_size)
-        results.append(tf.gather(embs, idx_list[i]))
   return results
 
 
diff --git a/hybridbackend/tensorflow/common/slice_sum.cu.cc b/hybridbackend/tensorflow/common/slice_sum.cu.cc
deleted file mode 100644
index 29726690..00000000
--- a/hybridbackend/tensorflow/common/slice_sum.cu.cc
+++ /dev/null
@@ -1,103 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#if GOOGLE_CUDA
-#define EIGEN_USE_GPU
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-#include <limits>
-
-#include <tensorflow/core/framework/register_types.h>
-#include <tensorflow/core/framework/tensor.h>
-#include <tensorflow/core/public/version.h>
-
-#include "hybridbackend/common/atomic.cu.h"
-#include "hybridbackend/tensorflow/common/device_functions.h"
-#include "hybridbackend/tensorflow/common/slice_sum.h"
-
-namespace tensorflow {
-
-using GPUDevice = Eigen::GpuDevice;
-
-namespace hybridbackend {
-
-namespace functor {
-
-template <typename T, int32 N = 256>
-__global__ void SliceSumKernel(const int32 num_rows, const int32 num_cols,
-                               const int32 col, const T* input, T* output_total,
-                               T* output) {
-  for (int32 idx : CudaGridRangeX(num_rows)) {
-    const T v = input[idx * num_cols + col];
-    output[idx] = v;
-    atomicAdd(output_total, v);
-  }
-}
-
-template <typename T>
-struct SliceSum<GPUDevice, T> {
-  void operator()(const int32 num_rows, const int32 num_cols, const int32 col,
-                  const T* input, T* output_total, T* output,
-                  const Eigen::GpuDevice& d) {
-    CudaLaunch(SliceSumKernel<T>, num_rows, 0, d, nullptr, num_rows, num_cols,
-               col, input, output_total, output);
-  }
-};
-
-template struct SliceSum<GPUDevice, int32>;
-template struct SliceSum<GPUDevice, int64>;
-template struct SliceSum<GPUDevice, uint32>;
-template struct SliceSum<GPUDevice, uint64>;
-
-template <typename T, int32 N = 256>
-__global__ void GroupSliceSumKernel(const int32 num_rows, const int32 num_cols,
-                                    const int32 col, const int32 num_inputs,
-                                    const T* inputs, T* output_totals,
-                                    T** outputs) {
-  for (int32 idx : CudaGridRangeX(num_inputs * num_rows)) {
-    const int32 s = idx / num_rows;
-    const int32 sidx = idx % num_rows;
-    const T v = inputs[idx * num_cols + col];
-    outputs[s][sidx] = v;
-    atomicAdd(output_totals + s, v);
-  }
-}
-
-template <typename T>
-struct SliceSumN<GPUDevice, T> {
-  void operator()(const int32 num_rows, const int32 num_cols, const int32 col,
-                  const int32 num_inputs, const T* inputs, T* output_totals,
-                  T** outputs, const Eigen::GpuDevice& d) {
-    CudaLaunch(GroupSliceSumKernel<T>, num_inputs * num_rows, 0, d, nullptr,
-               num_rows, num_cols, col, num_inputs, inputs, output_totals,
-               outputs);
-  }
-};
-
-template struct SliceSumN<GPUDevice, int32>;
-template struct SliceSumN<GPUDevice, int64>;
-template struct SliceSumN<GPUDevice, uint32>;
-template struct SliceSumN<GPUDevice, uint64>;
-
-}  // namespace functor
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // GOOGLE_CUDA
-#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/data/iterators.py b/hybridbackend/tensorflow/data/iterators.py
index 1612d833..a09948ec 100644
--- a/hybridbackend/tensorflow/data/iterators.py
+++ b/hybridbackend/tensorflow/data/iterators.py
@@ -31,9 +31,8 @@
 
 from hybridbackend.tensorflow.data.detect_end.dataset import DetectEndDataset
 from hybridbackend.tensorflow.data.prefetch.ops import Iterator
-from hybridbackend.tensorflow.distribute.communicator import CollectiveOps
-from hybridbackend.tensorflow.distribute.communicator_pool import \
-  CommunicatorPool
+from hybridbackend.tensorflow.distribute.collective import Collective
+from hybridbackend.tensorflow.distribute.ops import CollectiveOps
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.ops import ModeKeys
 from hybridbackend.tensorflow.framework.ops import MultiValues
@@ -168,12 +167,12 @@ def accept(cls, should_stop):
     with ops.device(Context.get().devices[Context.get().rank]):
       should_stop = math_ops.cast(should_stop, dtypes.int32)
       if Context.get().options.data_sync_drop_remainder:
-        should_stop_all = CommunicatorPool.get().allreduce(
-          should_stop, reduce_op=CollectiveOps.MAX, trainable=False)
+        should_stop_all = Collective.get().allreduce(
+          should_stop, reduce_op=CollectiveOps.MAX)
         return SessionRunRewriting.add_to_collection(
           DataSyncRewriting.__name__, should_stop_all)
-      should_stop_all = CommunicatorPool.get().allreduce(
-        should_stop, reduce_op=CollectiveOps.MIN, trainable=False)
+      should_stop_all = Collective.get().allreduce(
+        should_stop, reduce_op=CollectiveOps.MIN)
       return SessionRunRewriting.add_to_collection(
         DataSyncRewriting.__name__, should_stop_all)
 
diff --git a/hybridbackend/tensorflow/distribute/__init__.py b/hybridbackend/tensorflow/distribute/__init__.py
index 8b60fb16..b59cf9f1 100644
--- a/hybridbackend/tensorflow/distribute/__init__.py
+++ b/hybridbackend/tensorflow/distribute/__init__.py
@@ -20,12 +20,9 @@
 from __future__ import division
 from __future__ import print_function
 
-from hybridbackend.tensorflow.distribute.communicator import \
-  CollectiveOps as ops
-from hybridbackend.tensorflow.distribute.communicator import Communicator
-from hybridbackend.tensorflow.distribute.communicator_pool import \
-  CommunicatorPool
-from hybridbackend.tensorflow.distribute.nccl.comm import NcclCommunicator
+from hybridbackend.tensorflow.distribute.collective import *
+from hybridbackend.tensorflow.distribute.ops import CollectiveOps as ops
+from hybridbackend.tensorflow.distribute.partition.ops import *
 from hybridbackend.tensorflow.framework.context import Context as _ctx
 
 _ = (
@@ -33,4 +30,5 @@
   .register('comm_default', 'NCCL', env='HB_COMM_DEFAULT')
   .register('comm_pool_name', 'default')
   .register('comm_pool_capacity', 1)
-  .register('comm_pubsub_device', '/cpu:0'))
+  .register('comm_wire_dtype', None)
+  .register('comm_gradient_wire_dtype', None))
diff --git a/hybridbackend/tensorflow/distribute/allreduce_benchmark.py b/hybridbackend/tensorflow/distribute/allreduce_benchmark.py
deleted file mode 100644
index 87fe8220..00000000
--- a/hybridbackend/tensorflow/distribute/allreduce_benchmark.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''AllReduce benchmark.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import argparse
-
-import tensorflow as tf
-
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@hb.function()
-def build_bench_op(params):
-  comm_input = tf.get_variable(
-    'input',
-    initializer=tf.random_normal(
-      [params.kilobytes * 256],
-      mean=100,
-      stddev=80))
-  comm = hb.distribute.Communicator.build(
-    'benchmark', hb.context.devices)
-  comm_output = comm.allreduce(comm_input)
-  step = tf.train.get_or_create_global_step()
-  return tf.group([comm_output, step.assign_add(1)])
-
-
-def benchmark(params):
-  with hb.scope():
-    bench_op = build_bench_op(params)
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=[
-          tf.train.StopAtStepHook(params.num_steps),
-          hb.train.StepStatHook(
-            count=params.kilobytes / 1024.,
-            unit='MB')]) as sess:
-      while not sess.should_stop():
-        sess.run(bench_op)
-
-
-if __name__ == '__main__':
-  tf.logging.set_verbosity(tf.logging.INFO)
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--kilobytes', type=int, default=1024 * 64)  # 64MB
-  parser.add_argument('--num-steps', type=int, default=1000)
-  benchmark(parser.parse_args())
diff --git a/hybridbackend/tensorflow/distribute/alltoallw_benchmark.py b/hybridbackend/tensorflow/distribute/alltoallw_benchmark.py
deleted file mode 100644
index e5c03608..00000000
--- a/hybridbackend/tensorflow/distribute/alltoallw_benchmark.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''AllToAllw benchmark.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import argparse
-
-import tensorflow as tf
-
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@hb.function()
-def build_bench_op(params):
-  comm_inputs = [
-    tf.get_variable(
-      f'input{i}',
-      initializer=tf.random_normal(
-        [params.kilobytes * 256],
-        mean=100,
-        stddev=80))
-    for i, _ in enumerate(hb.context.devices)]
-  comm = hb.distribute.Communicator.build(
-    'benchmark', hb.context.devices)
-  comm_outputs = comm.alltoallw(comm_inputs)
-  step = tf.train.get_or_create_global_step()
-  return tf.group(comm_outputs + [step.assign_add(1)])
-
-
-def benchmark(params):
-  with hb.scope():
-    bench_op = build_bench_op(params)
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=[
-          tf.train.StopAtStepHook(params.num_steps),
-          hb.train.StepStatHook(
-            count=params.kilobytes / 1024.,
-            unit='MB')]) as sess:
-      while not sess.should_stop():
-        sess.run(bench_op)
-
-
-if __name__ == '__main__':
-  tf.logging.set_verbosity(tf.logging.INFO)
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--kilobytes', type=int, default=1024 * 64)  # 64MB
-  parser.add_argument('--num-steps', type=int, default=1000)
-  benchmark(parser.parse_args())
diff --git a/hybridbackend/tensorflow/distribute/alltoallw_n_benchmark.py b/hybridbackend/tensorflow/distribute/alltoallw_n_benchmark.py
deleted file mode 100644
index 0ef3f1bc..00000000
--- a/hybridbackend/tensorflow/distribute/alltoallw_n_benchmark.py
+++ /dev/null
@@ -1,69 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''AlltoallwN benchmark.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import argparse
-from six.moves import xrange  # pylint: disable=redefined-builtin
-
-import tensorflow as tf
-
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@hb.function()
-def build_bench_op(params):
-  comm_inputs_list = []
-  for i in xrange(params.num_columns):
-    comm_inputs_list.append([
-      tf.get_variable(
-        f'group{i}-input{j}',
-        shape=[params.kilobytes * 256],
-        initializer=tf.ones_initializer(tf.float32), dtype=tf.float32)
-      for j, _ in enumerate(hb.context.devices)])
-  comm = hb.distribute.Communicator.build(
-    'benchmark', hb.context.devices)
-  comm_outputs = comm.alltoallw_n(comm_inputs_list)
-  step = tf.train.get_or_create_global_step()
-  return tf.group(comm_outputs + [step.assign_add(1)])
-
-
-def benchmark(params):
-  with hb.scope():
-    bench_op = build_bench_op(params)
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=[
-          tf.train.StopAtStepHook(params.num_steps),
-          hb.train.StepStatHook(
-            count=params.num_columns * params.kilobytes / 1024.,
-            unit='MB')]) as sess:
-      while not sess.should_stop():
-        sess.run(bench_op)
-
-
-if __name__ == '__main__':
-  tf.logging.set_verbosity(tf.logging.INFO)
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--kilobytes', type=int, default=1024 * 64)  # 64MB
-  parser.add_argument('--num-steps', type=int, default=1000)
-  parser.add_argument('--num-columns', type=int, default=10)
-  benchmark(parser.parse_args())
diff --git a/hybridbackend/tensorflow/distribute/collective.h b/hybridbackend/tensorflow/distribute/collective.h
new file mode 100644
index 00000000..255d32c9
--- /dev/null
+++ b/hybridbackend/tensorflow/distribute/collective.h
@@ -0,0 +1,199 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_COLLECTIVE_H_
+#define HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_COLLECTIVE_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <tensorflow/core/framework/resource_mgr.h>
+#include <tensorflow/core/framework/tensor.h>
+#include <tensorflow/core/lib/core/errors.h>
+#include <tensorflow/core/public/version.h>
+
+#if GOOGLE_CUDA
+#include <tensorflow/core/common_runtime/gpu/gpu_event_mgr.h>
+#include <tensorflow/stream_executor/stream_executor.h>
+#endif
+
+#include "hybridbackend/tensorflow/common/stream.h"
+
+namespace tensorflow {
+
+#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) <= 1008L
+namespace se = ::perftools::gputools;
+#endif
+
+namespace hybridbackend {
+
+#if HYBRIDBACKEND_NCCL
+#if GOOGLE_CUDA
+
+enum CollectiveReduceOp {
+  kCollectiveSum = 0,
+  kCollectiveProd = 1,
+  kCollectiveMax = 2,
+  kCollectiveMin = 3,
+  kCollectiveAvg = 4
+};
+
+enum CollectiveTopology {
+  kCollectiveAllDevices = 0,
+  kCollectiveDevicesIntraNode = 1,
+  kCollectiveDevicesInterNode = 2
+};
+
+class Collective : public ResourceBase {
+ public:
+  string shared_name() const { return shared_name_; }
+
+  int world_size() const { return world_size_; }
+
+  int local_size() const { return local_size_; }
+
+  int rank() const { return rank_; }
+
+  void compute_active_ranks(CollectiveTopology topology,
+                            std::vector<int>& out) {
+    if (topology == kCollectiveDevicesIntraNode) {
+      int node_idx = rank_ / local_size_;
+      for (int rank = node_idx * local_size_;
+           rank < (node_idx + 1) * local_size_; ++rank) {
+        out.emplace_back(rank);
+      }
+    } else if (topology == kCollectiveDevicesInterNode) {
+      for (int rank = 0; rank < world_size_; ++rank) {
+        if (local_size_ == 1 || (rank % local_size_) == (rank_ % local_size_)) {
+          out.emplace_back(rank);
+        }
+      }
+    } else {
+      for (int rank = 0; rank < world_size_; ++rank) {
+        out.emplace_back(rank);
+      }
+    }
+  }
+
+  int compute_active_size(CollectiveTopology topology) {
+    int active_size;
+    if (topology == kCollectiveDevicesIntraNode) {
+      active_size = local_size_;
+    } else if (topology == kCollectiveDevicesInterNode) {
+      active_size = world_size_ / local_size_;
+    } else {
+      active_size = world_size_;
+    }
+
+    return active_size;
+  }
+
+#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) >= 1014L
+  string DebugString() const override {
+#else
+  string DebugString() override {
+#endif
+    return absl::StrCat("Collective(shared_name=", shared_name_,
+                        ", world_size=", world_size_,
+                        ", local_size=", local_size_, ", rank=", rank_, ")");
+  }
+
+  // Lifecycle management
+
+  Collective(const string& shared_name, int world_size, int local_size,
+             int rank)
+      : shared_name_(shared_name),
+        world_size_(world_size),
+        local_size_(local_size),
+        rank_(rank) {}
+
+  virtual ~Collective() { Destroy(); }
+
+  virtual Status Initialize(OpKernelContext* ctx) {
+    if (!TF_PREDICT_TRUE(0 <= rank_ && rank_ < world_size_)) {
+      return errors::InvalidArgument(absl::StrCat("Collective rank ", rank_,
+                                                  " or world_size ",
+                                                  world_size_, " is invalid."));
+    }
+
+    return Status::OK();
+  };
+
+  virtual Status Destroy() { return Status::OK(); };
+
+  // Symmetric collective operations
+
+  virtual Status Allreduce(const Tensor& input,
+                           const CollectiveReduceOp reduce_op,
+                           Tensor* output) = 0;
+
+  virtual Status AllreduceN(const std::vector<Tensor>& n_input,
+                            const CollectiveReduceOp reduce_op,
+                            std::vector<Tensor*>* n_output) = 0;
+
+  virtual Status Alltoall(const Tensor& input, Tensor* output,
+                          CollectiveTopology topology) = 0;
+
+  virtual Status AlltoallN(const std::vector<Tensor>& n_input,
+                           std::vector<Tensor*>* n_output,
+                           CollectiveTopology topology) = 0;
+
+  virtual Status AlltoallN(const std::vector<Tensor*>& n_input,
+                           std::vector<Tensor*>* n_output,
+                           CollectiveTopology topology) = 0;
+
+  virtual Status Alltoallv(const Tensor& input, const int32* send_sizes,
+                           const int32* recv_sizes, const int64 common_size,
+                           Tensor* output, CollectiveTopology topology) = 0;
+
+  virtual Status AlltoallvN(const std::vector<Tensor>& n_input,
+                            const std::vector<int32*>& n_send_sizes,
+                            const std::vector<int32*>& n_recv_sizes,
+                            const std::vector<int64>& n_common_size,
+                            std::vector<Tensor*>* n_output,
+                            CollectiveTopology topology) = 0;
+
+  virtual Status AlltoallvN(const std::vector<Tensor*>& n_input,
+                            const std::vector<int32*>& n_send_sizes,
+                            const std::vector<int32*>& n_recv_sizes,
+                            const std::vector<int64>& n_common_size,
+                            std::vector<Tensor*>* n_output,
+                            CollectiveTopology topology) = 0;
+
+  // Asymmetric collective operations
+
+  virtual Status Broadcast(const Tensor& input, const int root_rank,
+                           Tensor* output) = 0;
+
+  virtual Status Allgather(const Tensor& input, Tensor* output) = 0;
+
+  virtual Status Allgatherv(const Tensor& input, const Tensor& host_sizes,
+                            Tensor* output) = 0;
+
+ private:
+  string shared_name_;
+  int world_size_;
+  int local_size_;
+  int rank_;
+};
+
+#endif
+#endif
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif
+
+#endif  // HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_COLLECTIVE_H_
diff --git a/hybridbackend/tensorflow/distribute/collective.py b/hybridbackend/tensorflow/distribute/collective.py
new file mode 100644
index 00000000..fe99f3d4
--- /dev/null
+++ b/hybridbackend/tensorflow/distribute/collective.py
@@ -0,0 +1,549 @@
+# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+r'''Classes for collective communication.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import custom_gradient
+
+from hybridbackend.tensorflow.distribute.ops import CollectiveOps
+from hybridbackend.tensorflow.distribute.ops import Topology
+from hybridbackend.tensorflow.distribute.rpc import RpcCollective
+from hybridbackend.tensorflow.framework.context import Context
+from hybridbackend.tensorflow.framework.view import OperationLike
+
+
+class Collective(object):  # pylint: disable=useless-object-inheritance
+  r'''Resource for collective communication.
+  '''
+  def __init__(self, shared_name, handle, create_op, is_initialized_op):
+    r'''Constructs an instance for collective communication.
+
+    Args:
+      shared_name: shared name for collective communication.
+      handle: Handle to the resource.
+      create_op: Operation to create the resource.
+      is_initialized_op: Operation to check if the reousrce is initialized.
+    '''
+    self._shared_name = shared_name
+    self._handle = handle
+    self._create_op = create_op
+    self._is_initialized_op = is_initialized_op
+
+  @classmethod
+  def _resource_handle(cls, shared_name):
+    r'''Resource handle op.
+    '''
+    return (
+      OperationLike('CollectiveHandleOp')
+      .returns_tensor({}, dtypes.resource)
+      .finalize(shared_name=shared_name))
+
+  @classmethod
+  def _resource_get_id(cls, dtype, shape):
+    r'''Resource get_id op.
+    '''
+    return (
+      OperationLike('GetCollectiveId').returns_tensor(shape, dtype).finalize())
+
+  @classmethod
+  def _resource_create(cls, handle, collective_id, shared_name):
+    r'''Resource get_id op.
+    '''
+    return (
+      OperationLike('CreateCollective')
+      .finalize(
+        handle, collective_id,
+        world_size=Context.get().world_size,
+        local_size=Context.get().local_world_size,
+        rank=Context.get().rank,
+        shared_name=shared_name))
+
+  @classmethod
+  def _resource_is_initialized(cls, handle):
+    r'''Resource is_initialized op.
+    '''
+    return (
+      OperationLike('IsCollectiveInitialized')
+      .returns_tensor({}, dtypes.bool)
+      .finalize(handle))
+
+  @classmethod
+  def get(cls):
+    r'''Get or create an instance.
+    '''
+    ctx = Context.get()
+    if ctx.world_size == 1:
+      return cls('', None, None, None)
+
+    collective_resources = ops.get_collection_ref(cls.__name__)
+    if collective_resources:
+      return collective_resources[0]
+
+    shared_name = ops.get_default_graph().unique_name('collective')
+    with ops.name_scope(f'{shared_name}/replicas/{ctx.rank}/'):
+      with ops.control_dependencies(None):
+        handle = cls._resource_handle(shared_name)
+        collective_id_dtype = dtypes.int64
+        collective_id_shape = tensor_shape.TensorShape([16])  # 128 / 8
+        collective_id = (
+          RpcCollective(ctx.devices, ctx.rank).broadcast(
+            lambda: cls._resource_get_id(
+              collective_id_dtype, collective_id_shape),
+            collective_id_dtype, collective_id_shape,
+            name=f'{shared_name}_id/rpc_broadcast'))
+        create_op = cls._resource_create(handle, collective_id, shared_name)
+        is_initialized_op = cls._resource_is_initialized(handle)
+        coll = cls(shared_name, handle, create_op, is_initialized_op)
+        ops.add_to_collection(ops.GraphKeys.LOCAL_RESOURCES, coll)
+        collective_resources.append(coll)
+        return coll
+
+  @property
+  def shared_name(self):
+    r'''Shared name of the communicator.
+    '''
+    return self._shared_name
+
+  @property
+  def name(self):
+    r'''Name of the communicator instance.
+    '''
+    return f'{self.shared_name}/replicas/{Context.get().rank}/'
+
+  @property
+  def handle(self):
+    r'''Handle op to the communicator instance.
+    '''
+    return self._handle.op
+
+  @property
+  def create(self):
+    r'''Resource creation op of the communicator.
+    '''
+    return self._create_op  # pylint: disable=protected-access
+
+  @property
+  def is_initialized(self):
+    r'''Resource creation check op of the communicator.
+    '''
+    return self._is_initialized_op  # pylint: disable=protected-access
+
+  # All to all operations
+
+  def _allreduce(self, value, reduce_op, name):
+    r'''Reduce values across devices.
+
+    Args:
+      value: Value to be allreduced.
+      reduce_op: Reduction ops: SUM, PROD, MAX or MIN.
+      name: Name of the op.
+
+    Returns:
+      Allreduced value.
+    '''
+    if Context.get().world_size == 1:
+      return value
+
+    with ops.device(value.device):
+      with ops.name_scope(self.name):
+        return (
+          OperationLike('CollectiveAllreduce')
+          .returns_tensor(value.shape, value.dtype)
+          .finalize(self._handle, value, reduce_op=reduce_op, name=name))
+
+  def allreduce(self, value, reduce_op=None, name=None):
+    r'''Reduce values across devices.
+
+    Args:
+      value: Value to be allreduced.
+      reduce_op: Reduction ops: SUM, PROD, MAX or MIN.
+      name: Name of the op.
+
+    Returns:
+      Allreduced value.
+    '''
+    if reduce_op is None:
+      reduce_op = CollectiveOps.SUM
+    if name is None:
+      name = 'allreduce'
+
+    @custom_gradient.custom_gradient
+    def allreduce_fn(value):
+      reduced = self._allreduce(
+        value, reduce_op=reduce_op, name=name)
+
+      def grad_fn(value_grad):
+        r'''Gradient function for fn.
+        '''
+        if reduce_op != CollectiveOps.SUM:
+          raise NotImplementedError(
+            'Only reduce_op=SUM is supported for gradients computation.')
+        with ops.control_dependencies([reduced]):
+          grad_out = self._allreduce(
+            value_grad, reduce_op=reduce_op, name=f'gradients/{name}')
+        return grad_out
+      return reduced, grad_fn
+
+    return allreduce_fn(value)
+
+  def _alltoall(self, value, topology, wire_dtype, name):
+    r'''Shuffle value partitions across devices.
+
+    Args:
+      value: Value to be sent to other devices.
+      name: Name of the op.
+      topology: Communication topology.
+      wire_dtype: Data type for communication if possible.
+
+    Returns:
+      received value partitions from other devices.
+    '''
+    if Context.get().world_size == 1:
+      return value
+
+    with ops.device(value.device):
+      with ops.name_scope(self.name):
+        return (
+          OperationLike('CollectiveAlltoall')
+          .returns_tensor(value.shape, value.dtype)
+          .finalize(
+            self._handle, value,
+            wire_dtype=wire_dtype,
+            topology=topology,
+            name=name))
+
+  def _alltoallv(
+      self, value, sizes, common_shape,
+      topology, wire_dtype, name):
+    r'''Shuffle value partitions across devices.
+
+    Args:
+      value: Value to be sent to other devices.
+      sizes: Sizes of parts in the value.
+      common_shape: Common shape of values.
+      topology: Communication topology.
+      wire_dtype: Data type for communication if possible.
+
+    Returns:
+      received value partitions from other devices.
+    '''
+    if Context.get().world_size == 1:
+      return value
+
+    with ops.device(value.device):
+      with ops.name_scope(self.name):
+        return (
+          OperationLike('CollectiveAlltoallv')
+          .returns_tensors(
+            tensor_spec.TensorSpec(
+              tensor_shape.TensorShape([None]).concatenate(common_shape),
+              value.dtype),
+            tensor_spec.TensorSpec(sizes.shape, sizes.dtype))
+          .finalize(
+            self._handle, value, sizes,
+            wire_dtype=wire_dtype,
+            common_shape=tensor_shape.TensorShape(common_shape),
+            topology=topology,
+            name=name))
+
+  def alltoall(
+      self, value,
+      sizes=None,
+      common_shape=None,
+      topology=Topology.ALL,
+      name=None):
+    r'''Shuffle value partitions across devices.
+
+    Args:
+      value: Value to be sent to other devices.
+      sizes: Sizes of splits in corresponding value.
+      common_shape: common shape of tensors in value.
+      topology: Communication topology.
+      name: Name of the op.
+
+    Returns:
+      received value partitions from other devices.
+    '''
+    if common_shape is None:
+      common_shape = {}
+    wire_dtype = Context.get().options.comm_wire_dtype
+    if wire_dtype is None:
+      wire_dtype = dtypes.float32
+    gradient_wire_dtype = Context.get().options.comm_gradient_wire_dtype
+    if gradient_wire_dtype is None:
+      gradient_wire_dtype = wire_dtype
+    if name is None:
+      name = 'alltoall'
+
+    @custom_gradient.custom_gradient
+    def alltoall_fn(value):
+      exchanged = self._alltoall(
+        value,
+        topology=topology,
+        wire_dtype=wire_dtype,
+        name=name)
+
+      def grad_fn(value_grad):
+        r'''Gradient function for fn.
+        '''
+        value_grad = ops.convert_to_tensor(
+          value_grad, name=f'gradients/{name}/to_dense')
+        with ops.control_dependencies([exchanged]):
+          grad_out = self._alltoall(
+            value_grad,
+            topology=topology,
+            wire_dtype=gradient_wire_dtype,
+            name=f'gradients/{name}')
+        return grad_out
+      return exchanged, grad_fn
+
+    if sizes is None:
+      return alltoall_fn(value)
+
+    @custom_gradient.custom_gradient
+    def alltoallv_fn(value, sizes):
+      exchanged_value, exchanged_sizes = self._alltoallv(
+        value, sizes,
+        common_shape=tensor_shape.TensorShape(common_shape),
+        topology=topology,
+        wire_dtype=wire_dtype,
+        name=name)
+
+      def grad_fn(value_grad, sizes_grad):
+        r'''Gradient function for fn.
+        '''
+        del sizes_grad
+        value_grad = ops.convert_to_tensor(
+          value_grad, name=f'gradients/{name}/to_dense')
+        with ops.control_dependencies([exchanged_value, exchanged_sizes]):
+          value_grad_out, sizes_grad_out = self._alltoallv(
+            value_grad, exchanged_sizes,
+            common_shape=tensor_shape.TensorShape(common_shape),
+            topology=topology,
+            wire_dtype=gradient_wire_dtype,
+            name=f'gradients/{name}')
+        return value_grad_out, sizes_grad_out
+      return (exchanged_value, exchanged_sizes), grad_fn
+
+    return alltoallv_fn(value, sizes)
+
+  # One to all / all to all operations
+
+  def _broadcast(self, value, root_rank, name):
+    r'''Broadcast value across devices.
+
+    Args:
+      value: Value to broadcast.
+      root_rank: Rank of broadcast root.
+      name: Name of the op.
+
+    Returns:
+      Broadcasted value.
+    '''
+    if Context.get().world_size == 1:
+      return value
+
+    with ops.device(value.device):
+      with ops.name_scope(self.name):
+        return (
+          OperationLike('CollectiveBroadcast')
+          .returns_tensor(value.shape, value.dtype)
+          .finalize(self._handle, value, root_rank=root_rank, name=name))
+
+  def broadcast(self, value, root_rank=None, name=None):
+    r'''Broadcast value across devices.
+
+    Args:
+      value: Value to broadcast.
+      root_rank: Rank of broadcast root.
+      name: Name of the op.
+
+    Returns:
+      Broadcasted value.
+    '''
+    if root_rank is None:
+      root_rank = 0
+    if name is None:
+      name = 'broadcast'
+
+    return array_ops.stop_gradient(
+      self._broadcast(value, root_rank=root_rank, name=name))
+
+  def _allgather(self, value, name):
+    r'''Gather all values across devices to all devices.
+
+    Args:
+      value: Value on current device to be gathered.
+      varying_size: Supposing all value sizes on devices are not equal.
+      name: Name of the op.
+
+    Returns:
+      Gathered value.
+    '''
+    if Context.get().world_size == 1:
+      return value
+
+    with ops.device(value.device):
+      with ops.name_scope(self.name):
+        input_shape = value.shape
+        if input_shape.rank < 1:
+          input_shape = tensor_shape.TensorShape([1]).concatenate(input_shape)
+        output_shape = tensor_shape.TensorShape([
+          Context.get().world_size * input_shape.dims[0],
+          *input_shape.dims[1:]])
+        return (
+          OperationLike('CollectiveAllgather')
+          .returns_tensor(output_shape, value.dtype)
+          .finalize(self._handle, value, name=name))
+
+  def _allgatherv(self, value, name):
+    r'''Gather all values across devices to all devices.
+
+    Args:
+      value: Value on current device to be gathered.
+      varying_size: Supposing all value sizes on devices are not equal.
+      name: Name of the op.
+
+    Returns:
+      Gathered value.
+    '''
+    if Context.get().world_size == 1:
+      return value
+
+    with ops.device(value.device):
+      with ops.name_scope(self.name):
+        input_shape = value.shape
+        if input_shape.rank < 1:
+          input_shape = tensor_shape.TensorShape([1]).concatenate(input_shape)
+        output_shape = tensor_shape.TensorShape([None, *input_shape.dims[1:]])
+        return (
+          OperationLike('CollectiveAllgatherv')
+          .returns_tensor(output_shape, value.dtype)
+          .finalize(self._handle, value, name=name))
+
+  def allgather(self, value, varying_size=True, name=None):
+    r'''Gather all values across devices to all devices.
+
+    Args:
+      value: Value on current device to be gathered.
+      varying_size: Supposing all value sizes on devices are not equal.
+      name: Name of the op.
+
+    Returns:
+      Gathered value.
+    '''
+    if name is None:
+      name = 'allgather'
+
+    if varying_size:
+      return array_ops.stop_gradient(
+        self._allgatherv(value, name=name))
+
+    return array_ops.stop_gradient(
+      self._allgather(value, name=name))
+
+
+def active_size(topology):
+  r'''Obtain the size of active ranks from a topology.
+
+  Args:
+    topology: Communication topology.
+
+  Returns:
+    size of active ranks.
+  '''
+  if topology == Topology.INTRA_NODE:
+    return Context.get().local_world_size
+  if topology == Topology.INTER_NODE:
+    return int(
+      Context.get().world_size // Context.get().local_world_size)
+  return Context.get().world_size
+
+
+def allreduce(value, reduce_op=None, name=None):
+  r'''Reduce values across devices.
+
+  Args:
+    value: Value to be allreduced.
+    reduce_op: Reduction ops: SUM, PROD, MAX or MIN.
+    name: Name of the op.
+
+  Returns:
+    Allreduced value.
+  '''
+  return Collective.get().allreduce(
+    value,
+    reduce_op=reduce_op,
+    name=name)
+
+
+def alltoall(
+    value, sizes=None, common_shape=None, topology=Topology.ALL, name=None):
+  r'''Shuffle value partitions across devices.
+
+  Args:
+    value: Value to be sent to other devices.
+    sizes: Sizes of splits in corresponding value.
+    common_shape: common shape of tensors in value.
+    topology: Communication topology.
+    name: Name of the op.
+
+  Returns:
+    received value partitions from other devices.
+  '''
+  return Collective.get().alltoall(
+    value,
+    sizes=sizes,
+    common_shape=tensor_shape.TensorShape(common_shape),
+    topology=topology,
+    name=name)
+
+
+def broadcast(value, root_rank=None, name=None):
+  r'''Broadcast value across devices.
+
+  Args:
+    value: Value to broadcast.
+    root_rank: Rank of broadcast root.
+    name: Name of the op.
+
+  Returns:
+    Broadcasted value.
+  '''
+  return Collective.get().broadcast(value, root_rank=root_rank, name=name)
+
+
+def allgather(value, varying_size=True, name=None):
+  r'''Gather all values across devices to all devices.
+
+  Args:
+    value: Value on current device to be gathered.
+    varying_size: Supposing all value sizes on devices are not equal.
+    name: Name of the op.
+
+  Returns:
+    Gathered value.
+  '''
+  return Collective.get().allgather(value, varying_size=varying_size, name=name)
diff --git a/hybridbackend/tensorflow/distribute/communicator.py b/hybridbackend/tensorflow/distribute/communicator.py
deleted file mode 100644
index 9812a6aa..00000000
--- a/hybridbackend/tensorflow/distribute/communicator.py
+++ /dev/null
@@ -1,596 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Classes for collective communication.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import abc
-
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-
-from hybridbackend.tensorflow.framework.context import Context
-
-
-class CollectiveOps(object):  # pylint: disable=useless-object-inheritance
-  r'''Collective operations.
-  '''
-  SUM = 0
-  PROD = 1
-  MAX = 2
-  MIN = 3
-  AVG = 4
-
-
-class Communicator(object):  # pylint: disable=useless-object-inheritance
-  r'''A communicator for collective communication.
-  '''
-  class Resource(object):  # pylint: disable=useless-object-inheritance
-    r'''Resource object of a communciator.
-    '''
-    def __init__(self, comm):
-      self._comm = comm
-
-    @property
-    def name(self):
-      r'''Name of the communicator.
-      '''
-      return self._comm.name
-
-    @property
-    def handle(self):
-      r'''Resource handle of the communicator.
-      '''
-      return self._comm._handle.op  # pylint: disable=protected-access
-
-    @property
-    def create(self):
-      r'''Resource creation op of the communicator.
-      '''
-      return self._comm._create_op  # pylint: disable=protected-access
-
-    @property
-    def is_initialized(self):
-      r'''Resource creation check op of the communicator.
-      '''
-      return self._comm._is_initialized_op  # pylint: disable=protected-access
-
-  _registry = {}
-
-  @classmethod
-  def register(cls, impl):
-    cls._registry[impl.NAME] = impl
-
-  @classmethod
-  def as_type(cls, impl=None):
-    r'''Get communicator implementation type.
-    '''
-    if impl:
-      return impl
-    impl_type = Context.get().options.comm_default
-    if impl_type not in cls._registry:
-      raise ValueError(f'HB_COMM_DEFAULT is invalid: {impl_type}')
-    return cls._registry[impl_type]
-
-  @classmethod
-  def build(cls, shared_name, devices, impl=None, **kwargs):
-    r'''Builds a communicator.
-
-    Args:
-      shared_name: shared name of the communicator.
-      devices: devices of the communicator.
-      impl: implementation class for communication.
-      kwargs: (Optional.) key-value arguments.
-    '''
-    if not devices:
-      raise ValueError('devices must be provided')
-    devices = Context.canonicalize(devices)
-    if not devices:
-      raise ValueError('Devices must not be empty.')
-    return cls.as_type(impl)(shared_name, devices=devices, **kwargs)
-
-  def __init__(self, shared_name, devices, **kwargs):
-    r'''Constructs a communicator instance.
-
-    Args:
-      shared_name: shared name of the communicator.
-      devices: devices of the communicator.
-      kwargs: (Optional.) key-value arguments.
-    '''
-    if shared_name:
-      shared_name = shared_name.replace(':', '_').replace('/', '_')
-    else:
-      shared_name = ops.get_default_graph().unique_name('communicator')
-    self._shared_name = shared_name
-    self._device = Context.current_device()
-    if self._device not in devices:
-      raise ValueError(
-        f'Current device {self._device} not in devices {devices}')
-    self._devices = devices
-    self._kwargs = kwargs
-    if len(self._devices) > 1:
-      ops.add_to_collection(
-        ops.GraphKeys.LOCAL_RESOURCES,
-        self._build_resource())
-
-  @abc.abstractmethod
-  def _build_handle(self):
-    pass
-
-  @abc.abstractmethod
-  def _build_create_op(self):
-    pass
-
-  @abc.abstractmethod
-  def _build_is_initialized_op(self):
-    pass
-
-  def _build_resource(self):
-    r'''Create NcclCommunicator.Resource.
-    '''
-    with ops.name_scope(self.scope):
-      with ops.control_dependencies(None):
-        self._handle = self._build_handle()
-        self._create_op = self._build_create_op()
-        self._is_initialized_op = self._build_is_initialized_op()
-        return Communicator.Resource(self)
-
-  @property
-  def shared_name(self):
-    r'''Shared name of the communicator.
-    '''
-    return self._shared_name
-
-  @property
-  def name(self):
-    r'''Name of the communicator instance.
-    '''
-    return f'{self.shared_name}/replicas/{self.rank}'
-
-  @property
-  def device(self):
-    r'''Device of the communicator instance.
-    '''
-    return self._device
-
-  @property
-  def devices(self):
-    r'''All devices of the communicator.
-    '''
-    return self._devices
-
-  @property
-  def size(self):
-    r'''Size of the communicator.
-    '''
-    return len(self.devices)
-
-  @property
-  def rank(self):
-    r'''Rank of the communicator instance.
-    '''
-    return self.devices.index(self.device)
-
-  @property
-  def scope(self):
-    r'''Name scope of the communicator instance.
-    '''
-    return f'{self.shared_name}/{self.rank}'
-
-  @property
-  def handle(self):
-    r'''Handle to the communicator instance.
-    '''
-    return self._handle
-
-  # All-to-one computation
-
-  @abc.abstractmethod
-  def _reduce(self, value, reduce_op, root_rank, name):
-    pass
-
-  def reduce(self, value, reduce_op=CollectiveOps.SUM, root_rank=0, name=None):
-    r'''Reduce values across devices to the root device.
-
-    Args:
-      value: Value on current device to be reduced and scattered.
-      reduce_op: Reduction ops: SUM, PROD, MAX or MIN.
-      root_rank: Rank of broadcast root.
-      name: Name of the op.
-
-    Returns:
-      Reduced and scattered value on current device.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._reduce(value, reduce_op, root_rank, name)
-
-  # All-to-all computation
-
-  @abc.abstractmethod
-  def _reduce_scatter(self, value, reduce_op, name):
-    pass
-
-  def reduce_scatter(self, value, reduce_op=CollectiveOps.SUM, name=None):
-    r'''Reduce values across devices and scatter the result to all devices.
-
-    Args:
-      value: Value on current device to be reduced and scattered.
-      reduce_op: Reduction ops: SUM, PROD, MAX or MIN.
-      name: Name of the op.
-
-    Returns:
-      Reduced and scattered value on current device.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._reduce_scatter(value, reduce_op, name)
-
-  @abc.abstractmethod
-  def _allreduce(self, value, reduce_op, name):
-    pass
-
-  def allreduce(self, value, reduce_op=CollectiveOps.SUM, name=None):
-    r'''Reduce values across devices.
-
-    Args:
-      value: Value to be allreduced.
-      reduce_op: Reduction ops: SUM, PROD, MAX or MIN.
-      name: Name of the op.
-
-    Returns:
-      Allreduced value.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._allreduce(value, reduce_op, name)
-
-  # One-to-all data movement
-
-  @abc.abstractmethod
-  def _broadcast(self, value, root_rank, name):
-    pass
-
-  def broadcast(self, value, root_rank=0, name=None):
-    r'''Broadcast value across devices.
-
-    Args:
-      value: Value to broadcast.
-      root_rank: Rank of broadcast root.
-      name: Name of the op.
-
-    Returns:
-      Broadcasted value.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._broadcast(value, root_rank, name)
-
-  @abc.abstractmethod
-  def _scatter(self, value, root_rank, name):
-    pass
-
-  def scatter(self, value, root_rank=0, name=None):
-    r'''Scatter value on root device to all devices.
-
-    Args:
-      value: Value on current device to be scattered, no use for non-root.
-      root_rank: Rank of root scatter device.
-      name: Name of the op.
-
-    Returns:
-      Scattered value on current device.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._scatter(value, root_rank, name)
-
-  # All-to-one data movement
-
-  @abc.abstractmethod
-  def _gather(self, value, root_rank, name):
-    pass
-
-  def gather(self, value, root_rank=0, name=None):
-    r'''Gather all values across devices to root device.
-
-    Args:
-      value: Value on current device to be gathered.
-      root_rank: Rank of root gather device.
-      name: Name of the op.
-
-    Returns:
-      Gathered value on root, None or controlling object on non-root.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._gather(value, root_rank, name)
-
-  @abc.abstractmethod
-  def _gatherv(self, value, root_rank, name):
-    pass
-
-  def gatherv(self, value, root_rank=0, name=None):
-    r'''Gather all values with varying sizes across devices to root device.
-
-    Args:
-      value: Value on current device to be gathered.
-      root_rank: Rank of root gather device.
-      name: Name of the op.
-
-    Returns:
-      Gathered value on root, None or controlling object on non-root.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._gatherv(value, root_rank, name)
-
-  # All-to-all data movement
-
-  @abc.abstractmethod
-  def _allgather(self, value, name):
-    pass
-
-  def allgather(self, value, name=None):
-    r'''Gather all values across devices to all devices.
-
-    Args:
-      value: Value on current device to be gathered.
-      varying_size: Supposing all value sizes on devices are not equal.
-      name: Name of the op.
-
-    Returns:
-      Gathered value.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._allgather(value, name)
-
-  @abc.abstractmethod
-  def _allgatherv(self, value, name):
-    pass
-
-  def allgatherv(self, value, name=None):
-    r'''Gather all values across devices to all devices.
-
-    Args:
-      value: Value on current device to be gathered.
-      varying_size: Supposing all value sizes on devices are not equal.
-      name: Name of the op.
-
-    Returns:
-      Gathered value.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._allgatherv(value, name)
-
-  @abc.abstractmethod
-  def _alltoall(self, value, wire_dtype, name, **kwargs):
-    pass
-
-  def alltoall(self, value, wire_dtype=dtypes.float32, name=None, **kwargs):
-    r'''Shuffle value partitions across devices.
-
-    Args:
-      value: Value to be sent to other devices.
-      wire_dtype: Data type for communication if possible.
-      name: Name of the op.
-
-    Returns:
-      received value partitions from other devices.
-    '''
-    if self.size == 1:
-      return value
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._alltoall(value, wire_dtype, name, **kwargs)
-
-  @abc.abstractmethod
-  def _alltoallv(self, value, sizes, wire_dtype, common_shape, name, **kwargs):
-    pass
-
-  def alltoallv(
-      self, value, sizes,
-      wire_dtype=dtypes.float32,
-      common_shape=None,
-      name=None,
-      **kwargs):
-    r'''Shuffle value partitions with varying sizes across devices.
-
-    Args:
-      value: Value to be sent to other devices.
-      wire_dtype: Data type for communication if possible.
-      common_shape: common shape of tensors in value.
-      name: Name of the op.
-
-    Returns:
-      received value partitions from other devices.
-    '''
-    if self.size == 1:
-      return value, sizes
-
-    if common_shape is None:
-      common_shape = {}
-
-    with ops.name_scope(self.scope):
-      with ops.device(value.device):
-        return self._alltoallv(
-          value, sizes, wire_dtype, common_shape, name, **kwargs)
-
-  @abc.abstractmethod
-  def _alltoallv_n(
-      self, values, sizes, wire_dtype, common_shapes, name, **kwargs):
-    pass
-
-  def alltoallv_n(
-      self, values, sizes,
-      wire_dtype=dtypes.float32,
-      common_shapes=None,
-      name=None,
-      **kwargs):
-    r'''Grouped alltoallv
-
-    Args:
-      values: List of values to be sent to other devices.
-      sizes: List of sizes for values.
-      wire_dtype: Data type for communication if possible.
-      common_shapes: common shapes for values.
-      name: Name of the op.
-
-    Returns:
-      results: List of received values from other devices.
-      result_sizes: List of sizes for results.
-    '''
-    if self.size == 1:
-      return values, sizes
-
-    if not isinstance(values, (tuple, list)):
-      raise ValueError('values must be list')
-    values = list(values)
-    if not isinstance(sizes, (tuple, list)):
-      raise ValueError('sizes must be list')
-    sizes = list(sizes)
-    if len(values) != len(sizes):
-      raise ValueError('Length of values and length of sizes must be same')
-    if len({v.device for v in values + sizes}) > 1:
-      raise ValueError('all values and sizes must be placed at same device')
-    if not values:
-      return values, sizes
-
-    if common_shapes is None:
-      common_shapes = [{} for _ in values]
-    elif len(values) != len(common_shapes):
-      raise ValueError(
-        'Length of values and length of common_shapes must be same')
-
-    with ops.name_scope(self.scope):
-      with ops.device(values[0].device):
-        return self._alltoallv_n(
-          values, sizes, wire_dtype, common_shapes, name, **kwargs)
-
-  @abc.abstractmethod
-  def _alltoallw(self, values, wire_dtype, common_shape, name, **kwargs):
-    pass
-
-  def alltoallw(
-      self, values,
-      wire_dtype=dtypes.float32,
-      common_shape=None,
-      name=None,
-      **kwargs):
-    r'''Shuffle values with varying sizes and types across devices.
-
-    Args:
-      values: List of values to be sent to other devices.
-      wire_dtype: Data type for communication if possible.
-      common_shape: common shape of tensors in value.
-      name: Name of the op.
-
-    Returns:
-      received values from other devices.
-    '''
-    if self.size == 1:
-      return values
-
-    if common_shape is None:
-      common_shape = {}
-
-    if not isinstance(values, (list, tuple)):
-      raise ValueError('Only list or tuple allowed.')
-    if len(values) != self.size:
-      raise ValueError('Number of values must be same to number of devices')
-    if len({v.device for v in values}) > 1:
-      raise ValueError('Values must be placed at same device')
-    values_device = values[0].device
-
-    with ops.name_scope(self.scope):
-      with ops.device(values_device):
-        return self._alltoallw(
-          values, wire_dtype, common_shape, name, **kwargs)
-
-  @abc.abstractmethod
-  def _alltoallw_n(
-      self, values_list, wire_dtype, common_shapes, name, **kwargs):
-    pass
-
-  def alltoallw_n(
-      self, values_list,
-      wire_dtype=dtypes.float32,
-      common_shapes=None,
-      name=None,
-      **kwargs):
-    r'''Grouped alltoallw
-
-    Args:
-      values: List of values to be sent to other devices.
-      wire_dtype: Data type for communication if possible.
-      common_shape: common shape of tensors in value.
-      name: Name of the op.
-
-    Returns:
-      received values from other devices.
-    '''
-    if self.size == 1:
-      return values_list
-    if common_shapes is None:
-      common_shapes = [{} for _ in values_list]
-
-    if len({v.device for v in sum(values_list, [])}) > 1:
-      raise ValueError('all inputs must be placed at same device')
-    for values in values_list:
-      if len(values) != self.size:
-        raise ValueError('Number of inputs must be same to devices')
-    if common_shapes is None:
-      common_shapes = [{} for _ in values_list]
-    if not values_list:
-      return values_list
-
-    with ops.name_scope(self.scope):
-      with ops.device(values_list[0][0].device):
-        return self._alltoallw_n(
-          values_list, wire_dtype, common_shapes, name, **kwargs)
diff --git a/hybridbackend/tensorflow/distribute/communicator_pool.py b/hybridbackend/tensorflow/distribute/communicator_pool.py
deleted file mode 100644
index 2e093e82..00000000
--- a/hybridbackend/tensorflow/distribute/communicator_pool.py
+++ /dev/null
@@ -1,434 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Classes for collective communication pooling.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from six.moves import xrange  # pylint: disable=redefined-builtin
-
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import custom_gradient
-
-try:
-  from tensorflow.python.training import device_util
-except:  # pylint: disable=bare-except
-  from tensorflow.python.distribute import device_util
-
-from hybridbackend.tensorflow.distribute.communicator import CollectiveOps
-from hybridbackend.tensorflow.distribute.communicator import Communicator
-from hybridbackend.tensorflow.framework.context import Context
-
-
-class CommunicatorPool(object):  # pylint: disable=useless-object-inheritance
-  r'''Pool of communicators.
-  '''
-  _instances = {}
-
-  @classmethod
-  def get(cls, name=None):
-    r'''Get singleton with specific name.
-
-    Args:
-      name: Name of the communicator pool.
-    '''
-    name = name or Context.get().options.comm_pool_name
-    if name not in cls._instances:
-      cls._instances[name] = cls()
-    return cls._instances[name]
-
-  def __init__(self, impl=None, capacity=None):
-    r'''Initialize the communciator pool.
-
-    Args:
-      capacity: Capacity of the communciator pool.
-    '''
-    capacity = capacity or Context.get().options.comm_pool_capacity
-    self._impl = Communicator.as_type(impl)
-    self._capacity = self._impl.compute_pool_capacity(capacity)
-    self._pool = []
-    self._calls = 0
-    self._deps = []
-
-  @property
-  def capacity(self):
-    r'''Number of communicators in the pool.
-    '''
-    return self._capacity
-
-  @property
-  def size(self):
-    r'''Size of communicators in the pool.
-    '''
-    return len(self._pool)
-
-  def setup(self, impl=None, capacity=None):
-    r'''Set up the communicator pool.
-    Args:
-      impl: Implementation of communicator.
-      capacity: Number of communicators in the pool.
-    '''
-    if self._calls > 0:
-      raise ValueError('Pool cannot be set up after use')
-    self._impl = Communicator.as_type(impl)
-    self._capacity = self._impl.compute_pool_capacity(capacity)
-
-  def _build_communicator(self):
-    r'''Create a new communicator.
-    '''
-    comm_name = f'{Context.get().options.comm_pool_name}_comm'
-    if 'GPU' not in Context.get().current_device().upper():
-      local_gpus = [
-        device_util.canonicalize(f'/gpu:{d}')
-        for d in xrange(Context.get().num_gpus)]
-      with ops.device(f'/gpu:{Context.get().rank}'):
-        return Communicator.build(
-          comm_name,
-          local_gpus,
-          impl=self._impl)
-    return Communicator.build(
-      comm_name,
-      Context.get().devices,
-      impl=self._impl)
-
-  def call(self, fn, inputs, trainable=True, immutable=True):
-    r'''Call fn within communicator pool.
-
-    Args:
-      fn: Function accepts communicator, and returns results and grad function.
-      inputs: Inputs of fn.
-      trainable: If True, fn should return results and a function for gradients.
-      immutable: If True, grad(fn) do not accepts variables as arguments.
-
-    Returns:
-      results: Result of fn.
-    '''
-    if not isinstance(inputs, (tuple, list)):
-      inputs = [inputs]
-    inputs = list(inputs)
-    if self._capacity is None:
-      comm = self._build_communicator()
-      results, _ = fn(comm, inputs, None)
-      self._calls += 1
-      return results
-
-    if len(self._pool) < self._capacity:
-      self._pool.append(self._build_communicator())
-      self._deps.append(control_flow_ops.no_op())
-
-    if not trainable:
-      comm_idx = self._calls % self._capacity
-      comm = self._pool[comm_idx]
-      results, _ = fn(comm, inputs, [self._deps[comm_idx]])
-      self._deps[comm_idx] = control_flow_ops.group(results)
-      self._calls += 1
-      return results
-
-    @custom_gradient.custom_gradient
-    def wrapped_fn(*args):
-      r'''Call with custom gradients.
-      '''
-      args = list(args)
-      comm_idx = self._calls % self._capacity
-      comm = self._pool[comm_idx]
-      results, fn_grad = fn(comm, args, [self._deps[comm_idx]])
-      self._deps[comm_idx] = control_flow_ops.group(results)
-      self._calls += 1
-
-      if immutable:
-        def grad_fn(*grads):
-          r'''Gradient function for fn.
-          '''
-          d_inputs, _ = fn_grad(list(grads), [self._deps[comm_idx]], None)
-          self._deps[comm_idx] = control_flow_ops.group([
-            g.values if isinstance(g, ops.IndexedSlices) else g
-            for g in d_inputs if g is not None])
-          return d_inputs
-        return results, grad_fn
-
-      def mutable_grad_fn(*grads, **kwargs):
-        r'''Gradient function for fn.
-        '''
-        roots = list(kwargs.get('variables', []))
-        d_inputs, d_roots = fn_grad(
-          list(grads), [self._deps[comm_idx]], roots)
-        self._deps[comm_idx] = control_flow_ops.group([
-          g.values if isinstance(g, ops.IndexedSlices) else g
-          for g in d_inputs + d_roots if g is not None])
-        return d_inputs, d_roots
-      return results, mutable_grad_fn
-    return wrapped_fn(*inputs)
-
-  def reduce(
-      self, value,
-      reduce_op=CollectiveOps.SUM,
-      root_rank=0,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Reduce values across devices to the root device.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.reduce(
-          inputs[0],
-          reduce_op=reduce_op,
-          root_rank=root_rank,
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def reduce_scatter(
-      self, value,
-      reduce_op=CollectiveOps.SUM,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Reduce values across devices and scatter the result to all devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.reduce_scatter(
-          inputs[0],
-          reduce_op=reduce_op,
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def allreduce(
-      self, value,
-      reduce_op=CollectiveOps.SUM,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Reduce values across devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.allreduce(
-          inputs[0],
-          reduce_op=reduce_op,
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def broadcast(
-      self, value,
-      root_rank=0,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Broadcast value across devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.broadcast(
-          inputs[0],
-          root_rank=root_rank,
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def scatter(
-      self, value,
-      root_rank=0,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Scatter value on root device to all devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.scatter(
-          inputs[0],
-          root_rank=root_rank,
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def gather(
-      self, value,
-      root_rank=0,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Gather all values across devices to root device.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.gather(
-          inputs[0],
-          root_rank=root_rank,
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def gatherv(
-      self, value,
-      root_rank=0,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Gather all values with varying sizes across devices to root device.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.gatherv(
-          inputs[0],
-          root_rank=root_rank,
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def allgather(
-      self, value,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Gather all values across devices to all devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.allgather(
-          inputs[0],
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def allgatherv(
-      self, value,
-      trainable=True,
-      immutable=True,
-      name=None):
-    r'''Gather all values across devices to all devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.allgatherv(
-          inputs[0],
-          name=name), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def alltoall(
-      self, value,
-      wire_dtype=dtypes.float32,
-      trainable=True,
-      immutable=True,
-      name=None,
-      **kwargs):
-    r'''Shuffle value partitions across devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.alltoall(
-          inputs[0],
-          wire_dtype=wire_dtype,
-          name=name, **kwargs), None
-    return self.call(
-      _call_collective, value,
-      trainable=trainable,
-      immutable=immutable)
-
-  def alltoallv(
-      self, value, sizes,
-      wire_dtype=dtypes.float32,
-      common_shape=None,
-      trainable=True,
-      immutable=True,
-      name=None,
-      **kwargs):
-    r'''Shuffle value partitions with varying sizes across devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.alltoallv(
-          *inputs,
-          wire_dtype=wire_dtype,
-          common_shape=common_shape,
-          name=name, **kwargs), None
-    return self.call(
-      _call_collective, [value, sizes],
-      trainable=trainable,
-      immutable=immutable)
-
-  def alltoallw(
-      self, values,
-      wire_dtype=dtypes.float32,
-      common_shape=None,
-      trainable=True,
-      immutable=True,
-      name=None,
-      **kwargs):
-    r'''Shuffle values with varying sizes and types across devices.
-    '''
-    def _call_collective(comm, inputs, inputs_deps):
-      r'''Call collective function.
-      '''
-      with ops.control_dependencies(inputs_deps):
-        return comm.alltoallw(
-          inputs,
-          wire_dtype=wire_dtype,
-          common_shape=common_shape,
-          name=name, **kwargs), None
-    return self.call(
-      _call_collective, values,
-      trainable=trainable,
-      immutable=immutable)
diff --git a/hybridbackend/tensorflow/distribute/embedding.py b/hybridbackend/tensorflow/distribute/embedding.py
deleted file mode 100644
index 402355b0..00000000
--- a/hybridbackend/tensorflow/distribute/embedding.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Butterfly shuffle for embedding lookup.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import gradients
-
-from hybridbackend.tensorflow.distribute.communicator_pool import \
-  CommunicatorPool
-from hybridbackend.tensorflow.framework.context import Context
-from hybridbackend.tensorflow.ops.floormod_shuffle.ops import floormod_shuffle
-
-
-class ParallelEmbeddingLookupContext(object):  # pylint: disable=useless-object-inheritance
-  r'''Parallel embedding lookup context.
-  '''
-  def __init__(self, device):
-    self._shard_sizes = None
-    self._shard_unique_index = None
-    self._shard_index = None
-    self._device = device
-
-  @property
-  def device(self):
-    return self._device
-
-  def exchange_ids(self, ids):
-    r'''Exchange IDs across shards.
-    '''
-    num_shards = Context.get().world_size
-    ids_shards, ids_sizes, self._shard_index = floormod_shuffle(
-      ids, num_shards, name='shard_partition')
-    shard_ids, self._shard_sizes = CommunicatorPool.get().alltoallv(
-      ids_shards, ids_sizes, trainable=False)
-    shard_ids, self._shard_unique_index = array_ops.unique(
-      shard_ids, name='shard_unique')
-    return shard_ids
-
-  def exchange_embeddings(self, embeddings):
-    r'''Exchange embeddings across shards.
-    '''
-    def _comm_call_exchange_embeddings(dimension):
-      r'''Communicator call to exchange embeddings.
-      '''
-      def _fn(comm, inputs, inputs_deps):
-        r'''Communicator call to exchange embeddings.
-        '''
-        if inputs_deps is None:
-          embs, _ = comm.alltoallv(*inputs, common_shape=[dimension])
-          return embs, None
-        with ops.control_dependencies(inputs_deps):
-          embs, _ = comm.alltoallv(*inputs, common_shape=[dimension])
-
-        def grad_fn(grads, grads_deps, roots):
-          r'''Gradient function.
-          '''
-          del roots
-          with ops.control_dependencies(grads_deps):
-            d_inputs = gradients.gradients(embs, inputs, grad_ys=grads)
-          return d_inputs, None
-        return embs, grad_fn
-      return _fn
-    dimension = int(embeddings.shape[-1])
-    embeddings = array_ops.gather(
-      embeddings, self._shard_unique_index,
-      name='shard_unique_restore')
-    embeddings = CommunicatorPool.get().call(
-      _comm_call_exchange_embeddings(dimension),
-      [embeddings, self._shard_sizes])
-    embeddings = array_ops.gather(
-      embeddings, self._shard_index,
-      name='shard_stitch')
-    return embeddings
diff --git a/hybridbackend/tensorflow/distribute/nccl/collective.h b/hybridbackend/tensorflow/distribute/nccl/collective.h
new file mode 100644
index 00000000..63768330
--- /dev/null
+++ b/hybridbackend/tensorflow/distribute/nccl/collective.h
@@ -0,0 +1,136 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_NCCL_COLLECTIVE_H_
+#define HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_NCCL_COLLECTIVE_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <tensorflow/core/framework/resource_mgr.h>
+#include <tensorflow/core/framework/tensor.h>
+#include <tensorflow/core/lib/core/errors.h>
+#include <tensorflow/core/public/version.h>
+
+#if GOOGLE_CUDA
+#include <tensorflow/core/common_runtime/gpu/gpu_event_mgr.h>
+#include <tensorflow/stream_executor/stream_executor.h>
+#endif
+
+#include "hybridbackend/tensorflow/common/stream.h"
+#include "hybridbackend/tensorflow/distribute/collective.h"
+#include "hybridbackend/tensorflow/distribute/nccl/types.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+#if HYBRIDBACKEND_NCCL
+#if GOOGLE_CUDA
+
+class NcclCollective : public Collective {
+ public:
+  Stream* stream() const { return stream_; }
+
+  // Lifecycle management
+  NcclCollective(const string& shared_name, int world_size, int local_size,
+                 int rank);
+
+  virtual Status Initialize(OpKernelContext* ctx) override;
+
+  virtual Status Destroy() override;
+
+  // Symmetric collective operations
+
+  virtual Status Allreduce(const Tensor& input,
+                           const CollectiveReduceOp reduce_op,
+                           Tensor* output) override;
+
+  virtual Status AllreduceN(const std::vector<Tensor>& n_input,
+                            const CollectiveReduceOp reduce_op,
+                            std::vector<Tensor*>* n_output) override;
+
+  virtual Status Alltoall(const Tensor& input, Tensor* output,
+                          CollectiveTopology topology) override;
+
+  virtual Status AlltoallN(const std::vector<Tensor>& n_input,
+                           std::vector<Tensor*>* n_output,
+                           CollectiveTopology topology) override;
+
+  virtual Status AlltoallN(const std::vector<Tensor*>& n_input,
+                           std::vector<Tensor*>* n_output,
+                           CollectiveTopology topology) override;
+
+  virtual Status Alltoallv(const Tensor& input, const int32* send_sizes,
+                           const int32* recv_sizes, const int64 common_size,
+                           Tensor* output,
+                           CollectiveTopology topology) override;
+
+  virtual Status AlltoallvN(const std::vector<Tensor>& n_input,
+                            const std::vector<int32*>& n_send_sizes,
+                            const std::vector<int32*>& n_recv_sizes,
+                            const std::vector<int64>& n_common_size,
+                            std::vector<Tensor*>* n_output,
+                            CollectiveTopology topology) override;
+
+  virtual Status AlltoallvN(const std::vector<Tensor*>& n_input,
+                            const std::vector<int32*>& n_send_sizes,
+                            const std::vector<int32*>& n_recv_sizes,
+                            const std::vector<int64>& n_common_size,
+                            std::vector<Tensor*>* n_output,
+                            CollectiveTopology topology) override;
+
+  // Asymmetric collective operations
+
+  virtual Status Broadcast(const Tensor& input, const int root_rank,
+                           Tensor* output) override;
+
+  virtual Status Allgather(const Tensor& input, Tensor* output) override;
+
+  virtual Status Allgatherv(const Tensor& input, const Tensor& host_sizes,
+                            Tensor* output) override;
+
+  // NCCL specific operations
+
+  virtual Status Create(const string& id);
+
+  virtual Status CheckAsyncErrors();
+
+ private:
+  ncclComm_t comm_;
+  bool created_;
+  Stream* stream_;
+  std::mutex mu_;
+};
+
+class NcclCollectiveAsyncOp : public AsyncOpKernel {
+ public:
+  explicit NcclCollectiveAsyncOp(OpKernelConstruction* ctx);
+
+  virtual void CollectiveComputeAsync(NcclCollective* coll,
+                                      OpKernelContext* ctx,
+                                      DoneCallback done) = 0;
+
+  void ComputeAsync(OpKernelContext* ctx,
+                    AsyncOpKernel::DoneCallback done) override;
+};
+
+#endif
+#endif
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif
+
+#endif  // HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_NCCL_COLLECTIVE_H_
diff --git a/hybridbackend/tensorflow/distribute/nccl/comm.h b/hybridbackend/tensorflow/distribute/nccl/comm.h
deleted file mode 100644
index 1f7930c4..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/comm.h
+++ /dev/null
@@ -1,186 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#ifndef HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_NCCL_COMM_H_
-#define HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_NCCL_COMM_H_
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#include <tensorflow/core/framework/resource_mgr.h>
-#include <tensorflow/core/framework/tensor.h>
-#include <tensorflow/core/lib/core/errors.h>
-#include <tensorflow/core/public/version.h>
-
-#if GOOGLE_CUDA
-#include <tensorflow/core/common_runtime/gpu/gpu_event_mgr.h>
-#include <tensorflow/stream_executor/stream_executor.h>
-#endif
-
-#include "hybridbackend/tensorflow/distribute/nccl/types.h"
-
-namespace tensorflow {
-
-#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) <= 1008L
-namespace se = ::perftools::gputools;
-#endif
-
-namespace hybridbackend {
-
-#if HYBRIDBACKEND_NCCL
-#if GOOGLE_CUDA
-
-class NcclComm : public ResourceBase {
- public:
-  int rank() const { return rank_; }
-
-  int size() const { return size_; }
-
-  se::Stream* ctx_stream() const { return ctx_stream_; }
-  cudaStream_t* stream() const { return stream_; }
-
-#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) >= 1014L
-  string DebugString() const override {
-#else
-  string DebugString() override {
-#endif
-    return debug_string_;
-  }
-
-  NcclComm();
-
-  ~NcclComm();
-
-  Status Initialize(int size, int rank, const string& shared_name,
-                    OpKernelContext* ctx);
-
-  Status Create(const string& id);
-
-  Status Destroy();
-
-  Status Abort();
-
-  Status UserRank(int* rank);
-
-  void RunAsync(const string& message, OpKernelContext* ctx,
-                AsyncOpKernel::DoneCallback done, std::function<void()> func);
-
-  se::Event* ThenRecordEvent();
-
-  void ThenWaitFor(se::Event* ev);
-
-  void BlockHostUntilDone();
-
-  // All-to-one computation
-
-  Status Reduce(const Tensor& input, const ncclRedOp_t reduce_op,
-                const int root_rank, Tensor* output);
-
-  // All-to-all computation
-
-  Status ReduceScatter(const Tensor& input, const ncclRedOp_t reduce_op,
-                       Tensor* output);
-
-  Status Allreduce(const Tensor& input, const ncclRedOp_t reduce_op,
-                   Tensor* output);
-
-  // One-to-all data movement
-
-  Status Broadcast(const Tensor& input, const int root_rank, Tensor* output);
-
-  // Scatter not supported.
-
-  // All-to-one data movement
-
-  // Gather not supported.
-  // Gatherv not supported.
-
-  // All-to-all data movement
-
-  Status Allgather(const Tensor& input, Tensor* output);
-
-  Status GroupAllgather(const std::vector<Tensor>& inputs, Tensor* output);
-
-  Status Allgatherv(const Tensor& input, const Tensor& host_sizes,
-                    Tensor* output);
-
-  Status Alltoall(const Tensor& input, Tensor* output);
-
-  Status Alltoallv(const Tensor& input, const Tensor& host_sizes,
-                   const int64 common_size, Tensor* output);
-
-  Status AlltoallvN(const std::vector<Tensor>& inputs, const Tensor& host_sizes,
-                    const std::vector<int64>& common_sizes,
-                    std::vector<Tensor*>* outputs);
-
-  Status AlltoallvN(const std::vector<Tensor*>& inputs,
-                    const Tensor& host_sizes,
-                    const std::vector<int64>& common_sizes,
-                    std::vector<Tensor*>* outputs);
-
-  Status Alltoallw(const std::vector<Tensor>& inputs,
-                   std::vector<Tensor*>* outputs);
-
-  Status Alltoallw(const std::vector<Tensor*>& inputs,
-                   std::vector<Tensor*>* outputs);
-
-  Status AlltoallwN(const std::vector<Tensor>& inputs, const int64 num_columns,
-                    std::vector<Tensor*>* outputs);
-
-  Status AlltoallwN(const std::vector<Tensor*>& inputs, const int64 num_columns,
-                    std::vector<Tensor*>* outputs);
-
- private:
-  ncclComm_t comm_;
-  int size_;
-  int rank_;
-  bool created_;
-  string debug_string_;
-  std::unique_ptr<thread::ThreadPool> thread_pool_;
-  se::Stream* ctx_stream_;
-  cudaStream_t* stream_;
-};
-
-class NcclCommAsyncOp : public AsyncOpKernel {
- public:
-  explicit NcclCommAsyncOp(OpKernelConstruction* ctx);
-
-  virtual void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                                    DoneCallback done) = 0;
-
-  void ComputeAsync(OpKernelContext* ctx,
-                    AsyncOpKernel::DoneCallback done) override;
-
-  se::Event* ThenRecordEvent(OpKernelContext* ctx);
-
-  void ThenWaitFor(OpKernelContext* ctx, se::Event* ev);
-
-  void ThenExecute(OpKernelContext* ctx, std::function<void()> func);
-
-  void ThenCopyToDevice(OpKernelContext* ctx, Tensor* dst, const Tensor& src);
-
-  void ThenCopyToHost(OpKernelContext* ctx, Tensor* dst, const Tensor& src);
-
-  void BlockHostUntilDone(OpKernelContext* ctx);
-};
-
-#endif
-#endif
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif
-
-#endif  // HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_NCCL_COMM_H_
diff --git a/hybridbackend/tensorflow/distribute/nccl/comm.py b/hybridbackend/tensorflow/distribute/nccl/comm.py
deleted file mode 100644
index 4a5b70a2..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/comm.py
+++ /dev/null
@@ -1,315 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''NCCL based collective commmunication.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.platform import tf_logging as logging
-
-from hybridbackend.tensorflow.common import oplib as _ops
-from hybridbackend.tensorflow.distribute.communicator import CollectiveOps
-from hybridbackend.tensorflow.distribute.communicator import Communicator
-from hybridbackend.tensorflow.distribute.pubsub import PubSub
-
-ops.NotDifferentiable('HbGetNcclId')
-
-ops.NotDifferentiable('HbNcclComm')
-ops.NotDifferentiable('HbCreateNcclComm')
-ops.NotDifferentiable('HbIsNcclCommInitialized')
-
-
-@ops.RegisterGradient('HbNcclReduce')
-def _reduce_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL reduce op.
-  '''
-  comm = op.inputs[0]
-  grad_in = args[0]
-  reduce_op = op.get_attr('reduce_op')
-  root_rank = op.get_attr('root_rank')
-  if reduce_op != CollectiveOps.SUM:
-    raise NotImplementedError(
-      'Only reduce_op=SUM is supported for gradients computation.')
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out = _ops.hb_nccl_broadcast(
-        comm, grad_in, root_rank=root_rank)
-  return None, grad_out
-
-
-@ops.RegisterGradient('HbNcclReduceScatter')
-def _reduce_scatter_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL reduce scatter op.
-  '''
-  comm = op.inputs[0]
-  grad_in = args[0]
-  reduce_op = op.get_attr('reduce_op')
-  if reduce_op != CollectiveOps.SUM:
-    raise NotImplementedError(
-      'Only reduce_op=SUM is supported for gradients computation.')
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out = _ops.hb_nccl_allgather(comm, grad_in)
-  return None, grad_out
-
-
-@ops.RegisterGradient('HbNcclAllreduce')
-def _allreduce_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL allreduce op.
-  '''
-  comm = op.inputs[0]
-  grad_in = args[0]
-  reduce_op = op.get_attr('reduce_op')
-  if reduce_op != CollectiveOps.SUM:
-    raise NotImplementedError(
-      'Only reduce_op=SUM is supported for gradients computation.')
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out = _ops.hb_nccl_allreduce(
-        comm, grad_in, reduce_op=reduce_op)
-  return None, grad_out
-
-
-ops.NotDifferentiable('HbNcclBroadcast')
-ops.NotDifferentiable('ScatterWithNcclComm')
-ops.NotDifferentiable('GatherWithNcclComm')
-ops.NotDifferentiable('GathervWithNcclComm')
-
-
-@ops.RegisterGradient('HbNcclAllgather')
-def _allgather_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL allgather op.
-  '''
-  comm = op.inputs[0]
-  grad_in = args[0]
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out = _ops.hb_nccl_reduce_scatter(
-        comm, grad_in, reduce_op=CollectiveOps.SUM)
-  return None, grad_out
-
-
-ops.NotDifferentiable('HbNcclAllgatherv')
-
-
-@ops.RegisterGradient('HbNcclAlltoall')
-def _alltoall_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL alltoall op.
-  '''
-  comm = op.inputs[0]
-  grad_in = args[0]
-  wire_dtype = op.get_attr('wire_dtype')
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out = _ops.hb_nccl_alltoall(
-        comm, grad_in,
-        wire_dtype=wire_dtype)
-  return None, grad_out
-
-
-@ops.RegisterGradient('HbNcclAlltoallv')
-def _nccl_alltoallv_grad(op, *args):
-  r'''Gradient for NCCL alltoallv op.
-  '''
-  comm = op.inputs[0]
-  grad_in = list(args)[0]
-  grad_sizes_in = op.outputs[1]
-  common_shape = op.get_attr('common_shape')
-  wire_dtype = op.get_attr('wire_dtype')
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out, grad_sizes_out = _ops.hb_nccl_alltoallv(
-        comm, grad_in, grad_sizes_in,
-        wire_dtype=wire_dtype,
-        common_shape=common_shape)
-  return None, grad_out, grad_sizes_out
-
-
-@ops.RegisterGradient('HbNcclAlltoallvN')
-def _alltoallv_n_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL alltoallv_n op.
-  '''
-  comm = op.inputs[0]
-  num_columns = op.get_attr('num_columns')
-  wire_dtype = op.get_attr('wire_dtype')
-  common_shapes = op.get_attr('common_shapes')
-  grad_in = args[:num_columns]
-  grad_sizes_in = op.outputs[num_columns:]
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out, _ = _ops.hb_nccl_alltoallv_n(
-        comm, grad_in, grad_sizes_in,
-        wire_dtype=wire_dtype,
-        common_shapes=common_shapes)
-  return (None, *grad_out, *[None for _ in range(num_columns)])
-
-
-@ops.RegisterGradient('HbNcclAlltoallw')
-def _alltoallw_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL alltoallw op.
-  '''
-  comm = op.inputs[0]
-  common_shape = op.get_attr('common_shape')
-  wire_dtype = op.get_attr('wire_dtype')
-  grad_in = list(args)
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out = _ops.hb_nccl_alltoallw(
-        comm, grad_in,
-        wire_dtype=wire_dtype,
-        common_shape=common_shape)
-  return [None] + grad_out
-
-
-@ops.RegisterGradient('HbNcclAlltoallwN')
-def _alltoallw_n_with_nccl_comm_grad(op, *args):
-  r'''Gradient for NCCL alltoallw_n op.
-  '''
-  comm = op.inputs[0]
-  num_columns = op.get_attr('num_columns')
-  wire_dtype = op.get_attr('wire_dtype')
-  common_shapes = op.get_attr('common_shapes')
-  grad_in = list(args)
-  with ops.device(op.device):
-    with ops.control_dependencies(op.outputs):
-      grad_out = _ops.hb_nccl_alltoallw_n(
-        comm, grad_in,
-        num_columns=num_columns,
-        wire_dtype=wire_dtype,
-        common_shapes=common_shapes)
-  return [None] + grad_out
-
-
-class NcclCommunicator(Communicator):
-  r'''A communicator using NCCL.
-  '''
-  NAME = 'NCCL'
-
-  @classmethod
-  def compute_pool_capacity(cls, capacity=None):
-    if capacity is not None and capacity != 1:
-      logging.warning('Multiple-communicators pooling is unsafe.')
-    return capacity or 1
-
-  def _build_handle(self):
-    return _ops.hb_nccl_comm_handle_op(shared_name=self.shared_name)
-
-  def _build_create_op(self):
-    nccl_id = PubSub(self.devices)(
-      _ops.hb_get_nccl_id,
-      tensor_shape.TensorShape([16]),  # 128 / 8
-      dtypes.int64,
-      name=f'{self.shared_name}/nccl_id')
-    return _ops.hb_create_nccl_comm(
-      self.handle, nccl_id,
-      size=self.size,
-      rank=self.rank,
-      shared_name=self.shared_name)
-
-  def _build_is_initialized_op(self):
-    return _ops.hb_is_nccl_comm_initialized(self.handle)
-
-  def _reduce(self, value, reduce_op, root_rank, name):
-    return _ops.hb_nccl_reduce(
-      self.handle, value,
-      reduce_op=reduce_op,
-      root_rank=root_rank,
-      name=name)
-
-  def _reduce_scatter(self, value, reduce_op, name):
-    return _ops.hb_nccl_reduce_scatter(
-      self.handle, value,
-      reduce_op=reduce_op,
-      name=name)
-
-  def _allreduce(self, value, reduce_op, name):
-    return _ops.hb_nccl_allreduce(
-      self.handle, value,
-      reduce_op=reduce_op,
-      name=name)
-
-  def _broadcast(self, value, root_rank, name):
-    return _ops.hb_nccl_broadcast(
-      self.handle, value,
-      root_rank=root_rank,
-      name=name)
-
-  def _scatter(self, value, root_rank, name):
-    raise NotImplementedError
-
-  def _gather(self, value, root_rank, name):
-    raise NotImplementedError
-
-  def _gatherv(self, value, root_rank, name):
-    raise NotImplementedError
-
-  def _allgather(self, value, name):
-    return _ops.hb_nccl_allgather(self.handle, value, name=name)
-
-  def _allgatherv(self, value, name):
-    return _ops.hb_nccl_allgatherv(self.handle, value, name=name)
-
-  def _alltoall(self, value, wire_dtype, name, **kwargs):
-    return _ops.hb_nccl_alltoall(
-      self.handle, value,
-      wire_dtype=wire_dtype,
-      name=name)
-
-  def _alltoallv(self, value, sizes, wire_dtype, common_shape, name, **kwargs):
-    return _ops.hb_nccl_alltoallv(
-      self.handle, value, sizes,
-      common_shape=common_shape,
-      wire_dtype=wire_dtype,
-      name=name)
-
-  def _alltoallv_n(
-      self, values, sizes, wire_dtype, common_shapes, name, **kwargs):
-    return _ops.hb_nccl_alltoallv_n(
-      self.handle, values, sizes,
-      common_shapes=common_shapes,
-      wire_dtype=wire_dtype,
-      name=name)
-
-  def _alltoallw(self, values, wire_dtype, common_shape, name, **kwargs):
-    return _ops.hb_nccl_alltoallw(
-      self.handle, values,
-      common_shape=common_shape,
-      wire_dtype=wire_dtype,
-      name=name)
-
-  def _alltoallw_n(
-      self, values_list, wire_dtype, common_shapes, name, **kwargs):
-    flatten_results = _ops.hb_nccl_alltoallw_n(
-      self._handle,
-      sum(values_list, []),
-      num_columns=len(values_list),
-      wire_dtype=wire_dtype,
-      common_shapes=common_shapes)
-    results = []
-    for v in flatten_results:
-      if results and len(results[-1]) < self.size:
-        results[-1].append(v)
-      else:
-        results.append([v])
-    return results
-
-
-Communicator.register(NcclCommunicator)
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_allgather.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_allgather.cc
index 4bfdf2bb..24b3a368 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_allgather.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_allgather.cc
@@ -21,7 +21,7 @@ limitations under the License.
 
 #include <vector>
 
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
 namespace hybridbackend {
@@ -29,10 +29,10 @@ namespace hybridbackend {
 #if HYBRIDBACKEND_NCCL
 
 REGISTER_OP("HbNcclAllgather")
-    .Output("output: T")
+    .Output("output: dtype")
     .Input("handle: resource")
-    .Input("input: T")
-    .Attr("T: {int8, uint8, int32, uint32, int64, uint64, half, float, double}")
+    .Input("input: dtype")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
     .SetIsStateful()
     .SetShapeFn([](shape_inference::InferenceContext* c) {
       shape_inference::ShapeHandle input = c->input(1);
@@ -62,38 +62,41 @@ input: A tensor to gather.
 )doc");
 
 #if GOOGLE_CUDA
-class NcclAllgatherOp : public NcclCommAsyncOp {
+class NcclAllgatherOp : public NcclCollectiveAsyncOp {
  public:
-  explicit NcclAllgatherOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {}
+  explicit NcclAllgatherOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {}
 
-  void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                            DoneCallback done) override {
+  void CollectiveComputeAsync(NcclCollective* coll, OpKernelContext* ctx,
+                              DoneCallback done) override {
     const Tensor* input;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
     // Allocate output with a N times larger buffer than input, here N is the
     // communicator size which indicates the device amounts.
     TensorShape out_shape(input->shape());
     if (out_shape.dims() == 0) {
-      out_shape.AddDim(comm->size());
+      out_shape.AddDim(coll->world_size());
     } else {
-      out_shape.set_dim(0, out_shape.dim_size(0) * comm->size());
+      out_shape.set_dim(0, out_shape.dim_size(0) * coll->world_size());
     }
     Tensor* output;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, out_shape, &output),
                          done);
 
-    comm->RunAsync(
-        "NcclAllgather", ctx, done, [input, output, this, comm, ctx, done]() {
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Allgather]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Allgather(*input, output), done);
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [input, output, this, coll, ctx, done]() {
+          VLOG(1) << coll->DebugString() << " [" << name() << "] [Allgather]";
+          OP_REQUIRES_OK_ASYNC(ctx, coll->Allgather(*input, output), done);
+          coll->stream()->BlockComputeUntilDone(ctx, done);
         });
   }
 };
 
-#define REGISTER_KERNEL(TYPE)                                               \
-  REGISTER_KERNEL_BUILDER(                                                  \
-      Name("HbNcclAllgather").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \
-      NcclAllgatherOp);
+#define REGISTER_KERNEL(TYPE)                                 \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAllgather")             \
+                              .Device(DEVICE_GPU)             \
+                              .TypeConstraint<TYPE>("dtype"), \
+                          NcclAllgatherOp);
 TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 #endif  // GOOGLE_CUDA
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_allgatherv.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_allgatherv.cc
index 1c10df82..b72aac77 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_allgatherv.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_allgatherv.cc
@@ -21,7 +21,7 @@ limitations under the License.
 
 #include <vector>
 
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
 namespace hybridbackend {
@@ -29,10 +29,10 @@ namespace hybridbackend {
 #if HYBRIDBACKEND_NCCL
 
 REGISTER_OP("HbNcclAllgatherv")
-    .Output("output: T")
+    .Output("output: dtype")
     .Input("handle: resource")
-    .Input("input: T")
-    .Attr("T: {int8, uint8, int32, uint32, int64, uint64, half, float, double}")
+    .Input("input: dtype")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
     .SetIsStateful()
     .SetShapeFn([](shape_inference::InferenceContext* c) {
       shape_inference::ShapeHandle input = c->input(1);
@@ -62,118 +62,121 @@ input: A tensor to gather.
 )doc");
 
 #if GOOGLE_CUDA
-class NcclAllgathervOp : public NcclCommAsyncOp {
+class NcclAllgathervOp : public NcclCollectiveAsyncOp {
  public:
-  explicit NcclAllgathervOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {}
+  explicit NcclAllgathervOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {}
 
-  void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                            DoneCallback done) override {
-    const Tensor* input;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
-
-    Tensor* host_sizes = new Tensor();
+  void CollectiveComputeAsync(NcclCollective* coll, OpKernelContext* ctx,
+                              DoneCallback done) override {
     AllocatorAttributes host_sizes_alloc_attrs;
     host_sizes_alloc_attrs.set_on_host(true);
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(DT_INT32, TensorShape({comm->size()}), host_sizes,
-                           host_sizes_alloc_attrs),
-        done);
 
+    const Tensor* input;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
+    Tensor* host_sizes = new Tensor();
     auto done_ = [host_sizes, done]() {
       delete host_sizes;
       done();
     };
-    comm->RunAsync(
-        "NcclAllgatherv", ctx, done_,
-        [input, host_sizes, this, comm, ctx, done_]() {
-          // Allocate scratch tensors.
-          Tensor host_input_sizes;
-          AllocatorAttributes input_sizes_alloc_attrs;
-          input_sizes_alloc_attrs.set_on_host(true);
-          input_sizes_alloc_attrs.set_gpu_compatible(true);
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT32, TensorShape({}), &host_input_sizes,
-                                 input_sizes_alloc_attrs),
-              done_);
-          host_input_sizes.scalar<int32>()() = input->NumElements();
-
-          Tensor* input_sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT32, TensorShape({}), input_sizes,
-                                 ctx->output_alloc_attr(0)),
-              done_);
-          ThenCopyToDevice(ctx, input_sizes, host_input_sizes);
-
-          Tensor* sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT32, TensorShape({comm->size()}), sizes,
-                                 ctx->output_alloc_attr(0)),
-              done_);
-          auto sizes_ready = ThenRecordEvent(ctx);
-
-          comm->ThenWaitFor(sizes_ready);
-          comm->BlockHostUntilDone();
-
-          // Collect sizes of all inputs across devices.
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Allgather]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Allgather(*input_sizes, sizes),
-                               done_);
-          auto allgather_ready = comm->ThenRecordEvent();
-
-          ThenWaitFor(ctx, allgather_ready);
-          ThenCopyToHost(ctx, host_sizes, *sizes);
-          BlockHostUntilDone(ctx);
-          delete input_sizes;
-          delete sizes;
-
-          // Now all Tensor sizes have been gathered, calculate the total size
-          // and allocate the output Tensor for communication.
-          int32 total_size = 0;
-          bool is_same_size = true;
-          for (int i = 0; i < comm->size(); ++i) {
-            int32 per_elements = host_sizes->flat<int32>()(i);
-            total_size += per_elements;
-            if (is_same_size && per_elements != input->NumElements()) {
-              is_same_size = false;
-            }
-          }
-          TensorShape out_shape(input->shape());
-          int32 sub_size = 1;
-          for (int i = 1; i < out_shape.dims(); ++i) {
-            sub_size *= out_shape.dim_size(i);
-          }
-          if (out_shape.dims() == 0) {
-            out_shape.AddDim(total_size);
-          } else {
-            out_shape.set_dim(0, total_size / sub_size);
-          }
-          Tensor* output;
-          OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, out_shape, &output),
-                               done_);
-          auto outputs_ready = ThenRecordEvent(ctx);
-
-          comm->ThenWaitFor(outputs_ready);
-          if (is_same_size) {
-            VLOG(1) << comm->DebugString() << " [" << name() << "] [Allgather]";
-            OP_REQUIRES_OK_ASYNC(ctx, comm->Allgather(*input, output), done_);
-          } else {
-            VLOG(1) << comm->DebugString() << " [" << name()
-                    << "] [Allgatherv]";
-            OP_REQUIRES_OK_ASYNC(
-                ctx, comm->Allgatherv(*input, *host_sizes, output), done_);
-          }
-        });
+
+    OP_REQUIRES_OK_ASYNC(
+        ctx,
+        ctx->allocate_temp(DT_INT32, TensorShape({coll->world_size()}),
+                           host_sizes, host_sizes_alloc_attrs),
+        done_);
+
+    coll->stream()->LaunchUntilComputeDone(ctx, [input, host_sizes, this, coll,
+                                                 ctx, done_]() {
+      // Allocate scratch tensors.
+      Tensor host_input_sizes;
+      AllocatorAttributes input_sizes_alloc_attrs;
+      input_sizes_alloc_attrs.set_on_host(true);
+      input_sizes_alloc_attrs.set_gpu_compatible(true);
+      OP_REQUIRES_OK_ASYNC(
+          ctx,
+          ctx->allocate_temp(DT_INT32, TensorShape({}), &host_input_sizes,
+                             input_sizes_alloc_attrs),
+          done_);
+      host_input_sizes.scalar<int32>()() = input->NumElements();
+
+      Tensor* input_sizes = new Tensor();
+      OP_REQUIRES_OK_ASYNC(
+          ctx,
+          ctx->allocate_temp(DT_INT32, TensorShape({}), input_sizes,
+                             ctx->output_alloc_attr(0)),
+          done_);
+      se::DeviceMemoryBase dst_ptr(
+          const_cast<char*>(input_sizes->tensor_data().data()),
+          input_sizes->TotalBytes());
+      ctx->op_device_context()->stream()->ThenMemcpy(
+          &dst_ptr, host_input_sizes.tensor_data().data(),
+          host_input_sizes.TotalBytes());
+
+      Tensor* sizes = new Tensor();
+      OP_REQUIRES_OK_ASYNC(
+          ctx,
+          ctx->allocate_temp(DT_INT32, TensorShape({coll->world_size()}), sizes,
+                             ctx->output_alloc_attr(0)),
+          done_);
+      coll->stream()->ThenWaitUntilComputeDone(ctx);
+      coll->stream()->BlockHostUntilDone();
+
+      // Collect sizes of all inputs across devices.
+      VLOG(1) << coll->DebugString() << " [" << name() << "] [Allgather]";
+      OP_REQUIRES_OK_ASYNC(ctx, coll->Allgather(*input_sizes, sizes), done_);
+      coll->stream()->ThenMemcpy(
+          const_cast<char*>(host_sizes->tensor_data().data()),
+          se::DeviceMemoryBase(const_cast<char*>(sizes->tensor_data().data()),
+                               sizes->TotalBytes()),
+          host_sizes->TotalBytes());
+      coll->stream()->BlockHostUntilDone();
+      delete input_sizes;
+      delete sizes;
+
+      // Now all Tensor sizes have been gathered, calculate the total size
+      // and allocate the output Tensor for communication.
+      int32 total_size = 0;
+      bool is_same_size = true;
+      for (int i = 0; i < coll->world_size(); ++i) {
+        int32 per_elements = host_sizes->flat<int32>()(i);
+        total_size += per_elements;
+        if (is_same_size && per_elements != input->NumElements()) {
+          is_same_size = false;
+        }
+      }
+      TensorShape out_shape(input->shape());
+      int32 sub_size = 1;
+      for (int i = 1; i < out_shape.dims(); ++i) {
+        sub_size *= out_shape.dim_size(i);
+      }
+      if (out_shape.dims() == 0) {
+        out_shape.AddDim(total_size);
+      } else {
+        out_shape.set_dim(0, total_size / sub_size);
+      }
+      Tensor* output;
+      OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, out_shape, &output),
+                           done_);
+      coll->stream()->ThenWaitUntilComputeDone(ctx);
+      if (is_same_size) {
+        VLOG(1) << coll->DebugString() << " [" << name() << "] [Allgather]";
+        OP_REQUIRES_OK_ASYNC(ctx, coll->Allgather(*input, output), done_);
+      } else {
+        VLOG(1) << coll->DebugString() << " [" << name() << "] [Allgatherv]";
+        OP_REQUIRES_OK_ASYNC(ctx, coll->Allgatherv(*input, *host_sizes, output),
+                             done_);
+      }
+      coll->stream()->BlockComputeUntilDone(ctx, done_);
+    });
   }
 };
 
-#define REGISTER_KERNEL(TYPE)                                                \
-  REGISTER_KERNEL_BUILDER(                                                   \
-      Name("HbNcclAllgatherv").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \
-      NcclAllgathervOp);
+#define REGISTER_KERNEL(TYPE)                                 \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAllgatherv")            \
+                              .Device(DEVICE_GPU)             \
+                              .TypeConstraint<TYPE>("dtype"), \
+                          NcclAllgathervOp);
 TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 #endif  // GOOGLE_CUDA
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_allreduce.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_allreduce.cc
index c480abbe..405f8ef3 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_allreduce.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_allreduce.cc
@@ -21,7 +21,7 @@ limitations under the License.
 
 #include <vector>
 
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
 namespace hybridbackend {
@@ -29,11 +29,11 @@ namespace hybridbackend {
 #if HYBRIDBACKEND_NCCL
 
 REGISTER_OP("HbNcclAllreduce")
-    .Output("output: T")
+    .Output("output: dtype")
     .Input("handle: resource")
-    .Input("input: T")
+    .Input("input: dtype")
     .Attr("reduce_op: int >= 0 = 0")
-    .Attr("T: {int8, uint8, int32, uint32, int64, uint64, half, float, double}")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
     .SetShapeFn([](shape_inference::InferenceContext* c) {
       c->set_output(0, c->input(1));
       return Status::OK();
@@ -49,40 +49,241 @@ reduce_op: Reduce ops: 0 for SUM, 1 for PROD, 2 for MAX, 3 for MIN.
 )doc");
 
 #if GOOGLE_CUDA
-class NcclAllreduceOp : public NcclCommAsyncOp {
+class NcclAllreduceOp : public NcclCollectiveAsyncOp {
  public:
-  explicit NcclAllreduceOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {
+  explicit NcclAllreduceOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
     int reduce_op;
     OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_op", &reduce_op));
-    OP_REQUIRES(ctx, reduce_op >= 0,
-                errors::InvalidArgument("reduce_op is invalid:", reduce_op));
-    OP_REQUIRES_OK(ctx, ReduceOpToNcclReduceOp(reduce_op, &reduce_op_));
+    reduce_op_ = static_cast<CollectiveReduceOp>(reduce_op);
   }
 
-  void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                            DoneCallback done) override {
+  void CollectiveComputeAsync(NcclCollective* coll, OpKernelContext* ctx,
+                              DoneCallback done) override {
     const Tensor* input;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
     Tensor* output;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, input->shape(), &output),
                          done);
 
-    comm->RunAsync(
-        "NcclAllreduce", ctx, done, [input, output, this, comm, ctx, done]() {
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Allreduce]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Allreduce(*input, reduce_op_, output),
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [input, output, this, coll, ctx, done]() {
+          VLOG(1) << coll->DebugString() << " [" << name() << "] [Allreduce]";
+          OP_REQUIRES_OK_ASYNC(ctx, coll->Allreduce(*input, reduce_op_, output),
                                done);
+          coll->stream()->BlockComputeUntilDone(ctx, done);
         });
   }
 
  private:
-  ncclRedOp_t reduce_op_;
+  CollectiveReduceOp reduce_op_;
 };
 
-#define REGISTER_KERNEL(TYPE)                                               \
-  REGISTER_KERNEL_BUILDER(                                                  \
-      Name("HbNcclAllreduce").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \
-      NcclAllreduceOp);
+#define REGISTER_KERNEL(TYPE)                                 \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAllreduce")             \
+                              .Device(DEVICE_GPU)             \
+                              .TypeConstraint<TYPE>("dtype"), \
+                          NcclAllreduceOp);
+TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
+#undef REGISTER_KERNEL
+#endif  // GOOGLE_CUDA
+
+REGISTER_OP("HbNcclAllreduceN")
+    .Output("n_output: N * dtype")
+    .Input("handle: resource")
+    .Input("n_input: N * dtype")
+    .Attr("reduce_op: int >= 0 = 0")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
+    .Attr("N: int >= 1 = 1")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      int64 N;
+      TF_RETURN_IF_ERROR(c->GetAttr("N", &N));
+      for (int64 n = 0; n < N; ++n) {
+        c->set_output(n, c->input(1 + n));
+      }
+      return Status::OK();
+    })
+    .SetIsStateful()
+    .Doc(R"doc(
+Packed Allreduce using a NCCL communicator.
+
+n_output: N reduced tensors.
+handle: Handle of a NCCL communicator.
+n_input: N tensors to reduce.
+reduce_op: Reduce ops: 0 for SUM, 1 for PROD, 2 for MAX, 3 for MIN.
+)doc");
+
+#if GOOGLE_CUDA
+class NcclAllreduceNOp : public NcclCollectiveAsyncOp {
+ public:
+  explicit NcclAllreduceNOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &N_));
+    int reduce_op;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_op", &reduce_op));
+    reduce_op_ = static_cast<CollectiveReduceOp>(reduce_op);
+  }
+
+  void CollectiveComputeAsync(NcclCollective* coll, OpKernelContext* ctx,
+                              DoneCallback done) override {
+    OpInputList n_input_list;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("n_input", &n_input_list), done);
+    std::vector<Tensor>* n_input = new std::vector<Tensor>;
+    std::vector<Tensor*>* n_output = new std::vector<Tensor*>(N_, nullptr);
+    auto done_ = [this, n_input, n_output, done]() {
+      delete n_input;
+      delete n_output;
+      done();
+    };
+
+    for (int64 idx = 0; idx < N_; ++idx) {
+      auto& input = n_input_list[idx];
+      n_input->emplace_back(input);
+      OP_REQUIRES_OK_ASYNC(
+          ctx, ctx->allocate_output(idx, input.shape(), &(n_output->at(idx))),
+          done_);
+    }
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [n_input, n_output, coll, ctx, done_, this]() {
+          VLOG(1) << coll->DebugString() << " [" << name() << "] [AllreduceN]";
+          OP_REQUIRES_OK_ASYNC(
+              ctx, coll->AllreduceN(*n_input, reduce_op_, n_output), done_);
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
+        });
+  }
+
+ private:
+  int64 N_;
+  CollectiveReduceOp reduce_op_;
+};
+
+#define REGISTER_KERNEL(TYPE)                                 \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAllreduceN")            \
+                              .Device(DEVICE_GPU)             \
+                              .TypeConstraint<TYPE>("dtype"), \
+                          NcclAllreduceNOp);
+TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
+#undef REGISTER_KERNEL
+#endif  // GOOGLE_CUDA
+
+REGISTER_OP("HbNcclAllreduceMergedN")
+    .Output("n_output: N * dtype")
+    .Input("handle: resource")
+    .Input("n_input: N * dtype")
+    .Attr("reduce_op: int >= 0 = 0")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
+    .Attr("N: int >= 1 = 1")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      int64 N;
+      TF_RETURN_IF_ERROR(c->GetAttr("N", &N));
+      for (int64 n = 0; n < N; ++n) {
+        c->set_output(n, c->input(1 + n));
+      }
+      return Status::OK();
+    })
+    .SetIsStateful()
+    .Doc(R"doc(
+Packed merged Allreduce using a NCCL communicator.
+
+n_output: N reduced tensors.
+handle: Handle of a NCCL communicator.
+n_input: N tensors to reduce.
+reduce_op: Reduce ops: 0 for SUM, 1 for PROD, 2 for MAX, 3 for MIN.
+)doc");
+
+#if GOOGLE_CUDA
+class NcclAllreduceMergedNOp : public NcclCollectiveAsyncOp {
+ public:
+  explicit NcclAllreduceMergedNOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &N_));
+    int reduce_op;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_op", &reduce_op));
+    reduce_op_ = static_cast<CollectiveReduceOp>(reduce_op);
+  }
+
+  void CollectiveComputeAsync(NcclCollective* coll, OpKernelContext* ctx,
+                              DoneCallback done) override {
+    std::vector<int64>* input_bytes_vec = new std::vector<int64>(N_, 0);
+    std::vector<Tensor*>* n_output = new std::vector<Tensor*>(N_, nullptr);
+    Tensor* buffer_input = new Tensor();
+    Tensor* buffer_output = new Tensor();
+    auto done_ = [input_bytes_vec, n_output, buffer_input, buffer_output,
+                  done]() {
+      delete input_bytes_vec;
+      delete n_output;
+      delete buffer_input;
+      delete buffer_output;
+      done();
+    };
+
+    OpInputList n_input_list;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("n_input", &n_input_list), done_);
+    int64 total_bytes = 0;
+    for (int idx = 0; idx < N_; ++idx) {
+      const auto input_bytes = n_input_list[idx].TotalBytes();
+      total_bytes += input_bytes;
+      input_bytes_vec->at(idx) = input_bytes;
+    }
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->allocate_temp(DT_INT8, {total_bytes}, buffer_input), done_);
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->allocate_temp(DT_INT8, {total_bytes}, buffer_output), done_);
+    int64 offset_bytes = 0;
+    for (int64 idx = 0; idx < N_; ++idx) {
+      OP_REQUIRES_OK_ASYNC(ctx,
+                           ctx->allocate_output(idx, n_input_list[idx].shape(),
+                                                &(n_output->at(idx))),
+                           done_);
+
+      se::DeviceMemoryBase dst_ptr(
+          const_cast<char*>(buffer_input->tensor_data().data()) + offset_bytes,
+          input_bytes_vec->at(idx));
+      ctx->op_device_context()->stream()->ThenMemcpy(
+          &dst_ptr,
+          se::DeviceMemoryBase(
+              const_cast<char*>(n_input_list[idx].tensor_data().data()),
+              input_bytes_vec->at(idx)),
+          input_bytes_vec->at(idx));
+      offset_bytes += input_bytes_vec->at(idx);
+    }
+
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [input_bytes_vec, n_output, buffer_input, buffer_output, coll, ctx,
+              done_, this]() {
+          VLOG(1) << coll->DebugString() << " [" << name()
+                  << "] [AllreduceMergedN]";
+          OP_REQUIRES_OK_ASYNC(
+              ctx, coll->Allreduce(*buffer_input, reduce_op_, buffer_output),
+              done_);
+          int64 offset_bytes = 0;
+          for (int idx = 0; idx < N_; ++idx) {
+            se::DeviceMemoryBase dst_ptr(
+                const_cast<char*>(n_output->at(idx)->tensor_data().data()),
+                input_bytes_vec->at(idx));
+            coll->stream()->ThenMemcpy(
+                &dst_ptr,
+                se::DeviceMemoryBase(
+                    const_cast<char*>(buffer_output->tensor_data().data()) +
+                        offset_bytes,
+                    input_bytes_vec->at(idx)),
+                input_bytes_vec->at(idx));
+            offset_bytes += input_bytes_vec->at(idx);
+          }
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
+        });
+  }
+
+ private:
+  int64 N_;
+  CollectiveReduceOp reduce_op_;
+};
+
+#define REGISTER_KERNEL(TYPE)                                 \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAllreduceMergedN")      \
+                              .Device(DEVICE_GPU)             \
+                              .TypeConstraint<TYPE>("dtype"), \
+                          NcclAllreduceMergedNOp);
 TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 #endif  // GOOGLE_CUDA
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoall.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_alltoall.cc
index e9d83ce9..1accb883 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoall.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_alltoall.cc
@@ -22,43 +22,24 @@ limitations under the License.
 #include <vector>
 
 #include "hybridbackend/tensorflow/common/cast.h"
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
 namespace hybridbackend {
 
 #if HYBRIDBACKEND_NCCL
 
-REGISTER_OP("HbNcclAlltoall")
-    .Output("output: dtype")
-    .Input("handle: resource")
-    .Input("input: dtype")
-    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
-    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
-    .SetIsStateful()
-    .SetShapeFn([](shape_inference::InferenceContext* c) {
-      c->set_output(0, c->input(1));
-      return Status::OK();
-    })
-    .Doc(R"doc(
-AllToAll using a NCCL communicator.
-
-output: Exchanged tensor for each device.
-handle: Handle of a NCCL communicator.
-input: Tensor to be exchanged bettween each device.
-)doc");
-
 #if GOOGLE_CUDA
-
 namespace functor {
 template <typename DTYPE, typename WIRE_DTYPE>
 struct NcclAlltoallCall {
   Status operator()(const Tensor& input, Tensor* output, Tensor* comm_input,
-                    Tensor* comm_output, OpKernelContext* ctx, NcclComm* comm,
-                    NcclCommAsyncOp* comm_op) {
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [Alltoall]";
-    TF_RETURN_IF_ERROR(comm->Alltoall(input, output));
+                    Tensor* comm_output, CollectiveTopology topology,
+                    OpKernelContext* ctx, NcclCollective* coll,
+                    NcclCollectiveAsyncOp* comm_op) {
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [Alltoall]";
+    TF_RETURN_IF_ERROR(coll->Alltoall(input, output, topology));
     return Status::OK();
   }
 };
@@ -66,38 +47,107 @@ struct NcclAlltoallCall {
 template <>
 struct NcclAlltoallCall<float, Eigen::half> {
   Status operator()(const Tensor& input, Tensor* output, Tensor* comm_input,
-                    Tensor* comm_output, OpKernelContext* ctx, NcclComm* comm,
-                    NcclCommAsyncOp* comm_op) {
+                    Tensor* comm_output, CollectiveTopology topology,
+                    OpKernelContext* ctx, NcclCollective* coll,
+                    NcclCollectiveAsyncOp* comm_op) {
     TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, input.shape(), comm_input));
     TF_RETURN_IF_ERROR(
         ctx->allocate_temp(DT_HALF, output->shape(), comm_output));
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastIn]";
-    functor::Cast<float, Eigen::half>()(input, comm_input, ctx, comm->stream());
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [Alltoall]";
-    TF_RETURN_IF_ERROR(comm->Alltoall(*comm_input, comm_output));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastOut]";
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastIn]";
+    functor::Cast<float, Eigen::half>()(input, comm_input, ctx,
+                                        coll->stream()->get());
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [Alltoall]";
+    TF_RETURN_IF_ERROR(coll->Alltoall(*comm_input, comm_output, topology));
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastOut]";
     functor::Cast<Eigen::half, float>()(*output, comm_output, ctx,
-                                        comm->stream());
+                                        coll->stream()->get());
+    return Status::OK();
+  }
+};
+
+template <typename DTYPE, typename WIRE_DTYPE>
+struct NcclAlltoallNCall {
+  Status operator()(const std::vector<Tensor>& n_input,
+                    std::vector<Tensor*>* n_output,
+                    std::vector<Tensor*>* n_comm_input,
+                    std::vector<Tensor*>* n_comm_output,
+                    CollectiveTopology topology, OpKernelContext* ctx,
+                    NcclCollective* coll, NcclCollectiveAsyncOp* comm_op) {
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name()
+            << "] [AlltoallN]";
+    TF_RETURN_IF_ERROR(coll->AlltoallN(n_input, n_output, topology));
     return Status::OK();
   }
 };
 
+template <>
+struct NcclAlltoallNCall<float, Eigen::half> {
+  Status operator()(const std::vector<Tensor>& n_input,
+                    std::vector<Tensor*>* n_output,
+                    std::vector<Tensor*>* n_comm_input,
+                    std::vector<Tensor*>* n_comm_output,
+                    CollectiveTopology topology, OpKernelContext* ctx,
+                    NcclCollective* coll, NcclCollectiveAsyncOp* comm_op) {
+    for (int idx = 0; idx < n_input.size(); ++idx) {
+      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, n_input[idx].shape(),
+                                            n_comm_input->at(idx),
+                                            ctx->input_alloc_attr(idx)));
+      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, n_output->at(idx)->shape(),
+                                            n_comm_output->at(idx),
+                                            ctx->output_alloc_attr(idx)));
+    }
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastIn]";
+    functor::CastN<float, Eigen::half>()(n_input, n_comm_input, ctx,
+                                         coll->stream()->get());
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name()
+            << "] [AlltoallN]";
+    TF_RETURN_IF_ERROR(coll->AlltoallN(*n_comm_input, n_comm_output, topology));
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastOut]";
+    functor::CastN<Eigen::half, float>()(*n_comm_output, n_output, ctx,
+                                         coll->stream()->get());
+    return Status::OK();
+  }
+};
 }  // namespace functor
+#endif  // GOOGLE_CUDA
 
+REGISTER_OP("HbNcclAlltoall")
+    .Output("output: dtype")
+    .Input("handle: resource")
+    .Input("input: dtype")
+    .Attr("topology: int = 0")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
+    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(1));
+      return Status::OK();
+    })
+    .Doc(R"doc(
+AllToAll using a NCCL communicator.
+
+output: Exchanged tensor.
+handle: Handle of a NCCL communicator.
+input: Tensor to exchange.
+)doc");
+
+#if GOOGLE_CUDA
 template <typename DTYPE, typename WIRE_DTYPE>
-class NcclAlltoallOp : public NcclCommAsyncOp {
+class NcclAlltoallOp : public NcclCollectiveAsyncOp {
  public:
-  explicit NcclAlltoallOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {}
-
-  virtual void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                                    DoneCallback done) override {
-    const Tensor* input;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
-    Tensor* output;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, input->shape(), &output),
-                         done);
+  explicit NcclAlltoallOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    int topology;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("topology", &topology));
+    topology_ = static_cast<CollectiveTopology>(topology);
+  }
 
+  virtual void CollectiveComputeAsync(NcclCollective* coll,
+                                      OpKernelContext* ctx,
+                                      DoneCallback done) override {
     Tensor* comm_input = new Tensor();
     Tensor* comm_output = new Tensor();
     auto done_ = [comm_input, comm_output, done]() {
@@ -106,24 +156,263 @@ class NcclAlltoallOp : public NcclCommAsyncOp {
       done();
     };
 
-    comm->RunAsync(
-        "NcclAlltoall", ctx, done_,
-        [input, output, comm_input, comm_output, ctx, comm, this, done_]() {
+    const Tensor* input;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done_);
+    Tensor* output;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, input->shape(), &output),
+                         done_);
+    coll->stream()->LaunchUntilComputeDone(
+        ctx,
+        [input, output, comm_input, comm_output, ctx, coll, this, done_]() {
+          auto call = functor::NcclAlltoallCall<DTYPE, WIRE_DTYPE>();
+          OP_REQUIRES_OK_ASYNC(ctx,
+                               call(*input, output, comm_input, comm_output,
+                                    topology_, ctx, coll, this),
+                               done_);
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
+        });
+  }
+
+ private:
+  CollectiveTopology topology_;
+};
+
+#define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoall")                         \
+                              .Device(DEVICE_GPU)                        \
+                              .TypeConstraint<DTYPE>("dtype")            \
+                              .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
+                          NcclAlltoallOp<DTYPE, WIRE_DTYPE>);
+TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
+#undef REGISTER_KERNEL
+#endif  // GOOGLE_CUDA
+
+REGISTER_OP("HbNcclAlltoallN")
+    .Output("n_output: N * dtype")
+    .Input("handle: resource")
+    .Input("n_input: N * dtype")
+    .Attr("topology: int = 0")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
+    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
+    .Attr("N: int >= 1 = 1")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      int64 N;
+      TF_RETURN_IF_ERROR(c->GetAttr("N", &N));
+      for (int64 n = 0; n < N; ++n) {
+        c->set_output(n, c->input(1 + n));
+      }
+      return Status::OK();
+    })
+    .Doc(R"doc(
+Packed AllToAll using a NCCL communicator.
+
+n_output: N exchanged tensors.
+handle: Handle of a NCCL communicator.
+n_input: N tensors to exchange.
+)doc");
+
+#if GOOGLE_CUDA
+template <typename DTYPE, typename WIRE_DTYPE>
+class NcclAlltoallNOp : public NcclCollectiveAsyncOp {
+ public:
+  explicit NcclAlltoallNOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &N_));
+    int topology;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("topology", &topology));
+    topology_ = static_cast<CollectiveTopology>(topology);
+  }
+
+  virtual void CollectiveComputeAsync(NcclCollective* coll,
+                                      OpKernelContext* ctx,
+                                      DoneCallback done) override {
+    std::vector<Tensor>* n_input = new std::vector<Tensor>();
+    std::vector<Tensor*>* n_comm_input = new std::vector<Tensor*>();
+    std::vector<Tensor*>* n_comm_output = new std::vector<Tensor*>();
+    for (int idx = 0; idx < N_; ++idx) {
+      n_comm_input->push_back(new Tensor());
+      n_comm_output->push_back(new Tensor());
+    }
+    std::vector<Tensor*>* n_output = new std::vector<Tensor*>();
+
+    auto done_ = [this, n_input, n_output, n_comm_input, n_comm_output,
+                  done]() {
+      delete n_input;
+      delete n_output;
+      for (int idx = 0; idx < N_; ++idx) {
+        delete n_comm_input->at(idx);
+        delete n_comm_output->at(idx);
+      }
+      delete n_comm_input;
+      delete n_comm_output;
+      done();
+    };
+
+    OpInputList n_input_list;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("n_input", &n_input_list), done_);
+    for (int idx = 0; idx < N_; ++idx) {
+      n_input->push_back(n_input_list[idx]);
+      Tensor* output;
+      OP_REQUIRES_OK_ASYNC(
+          ctx, ctx->allocate_output(idx, n_input_list[idx].shape(), &output),
+          done_);
+      n_output->push_back(output);
+    }
+
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [n_input, n_output, n_comm_input, n_comm_output, ctx, coll, this,
+              done_]() {
+          auto call = functor::NcclAlltoallNCall<DTYPE, WIRE_DTYPE>();
+          OP_REQUIRES_OK_ASYNC(ctx,
+                               call(*n_input, n_output, n_comm_input,
+                                    n_comm_output, topology_, ctx, coll, this),
+                               done_);
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
+        });
+  }
+
+ private:
+  CollectiveTopology topology_;
+  int64 N_;
+};
+
+#define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoallN")                        \
+                              .Device(DEVICE_GPU)                        \
+                              .TypeConstraint<DTYPE>("dtype")            \
+                              .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
+                          NcclAlltoallNOp<DTYPE, WIRE_DTYPE>);
+TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
+#undef REGISTER_KERNEL
+#endif  // GOOGLE_CUDA
+
+REGISTER_OP("HbNcclAlltoallMergedN")
+    .Output("n_output: N * dtype")
+    .Input("handle: resource")
+    .Input("n_input: N * dtype")
+    .Attr("topology: int = 0")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
+    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
+    .Attr("N: int >= 1 = 1")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      int64 N;
+      TF_RETURN_IF_ERROR(c->GetAttr("N", &N));
+      for (int64 n = 0; n < N; ++n) {
+        c->set_output(n, c->input(1 + n));
+      }
+      return Status::OK();
+    })
+    .Doc(R"doc(
+Packed merged AllToAll using a NCCL communicator.
+
+n_output: N exchanged tensors.
+handle: Handle of a NCCL communicator.
+n_input: N tensors to exchange.
+)doc");
+
+#if GOOGLE_CUDA
+template <typename DTYPE, typename WIRE_DTYPE>
+class NcclAlltoallMergedNOp : public NcclCollectiveAsyncOp {
+ public:
+  explicit NcclAlltoallMergedNOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &N_));
+    int topology;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("topology", &topology));
+    topology_ = static_cast<CollectiveTopology>(topology);
+  }
+
+  virtual void CollectiveComputeAsync(NcclCollective* coll,
+                                      OpKernelContext* ctx,
+                                      DoneCallback done) override {
+    std::vector<int64>* input_bytes_vec = new std::vector<int64>(N_, 0);
+    std::vector<Tensor*>* n_output = new std::vector<Tensor*>(N_, nullptr);
+    Tensor* buffer_input = new Tensor();
+    Tensor* buffer_output = new Tensor();
+    Tensor* buffer_comm_input = new Tensor();
+    Tensor* buffer_comm_output = new Tensor();
+    auto done_ = [input_bytes_vec, n_output, buffer_input, buffer_output,
+                  buffer_comm_input, buffer_comm_output, done]() {
+      delete input_bytes_vec;
+      delete n_output;
+      delete buffer_input;
+      delete buffer_output;
+      delete buffer_comm_input;
+      delete buffer_comm_output;
+      done();
+    };
+
+    OpInputList n_input_list;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("n_input", &n_input_list), done_);
+    int64 total_bytes = 0;
+    for (int idx = 0; idx < N_; ++idx) {
+      const auto input_bytes = n_input_list[idx].TotalBytes();
+      total_bytes += input_bytes;
+      input_bytes_vec->at(idx) = input_bytes;
+    }
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->allocate_temp(DT_INT8, {total_bytes}, buffer_input), done_);
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->allocate_temp(DT_INT8, {total_bytes}, buffer_output), done_);
+    int64 offset_bytes = 0;
+    for (int64 idx = 0; idx < N_; ++idx) {
+      OP_REQUIRES_OK_ASYNC(ctx,
+                           ctx->allocate_output(idx, n_input_list[idx].shape(),
+                                                &(n_output->at(idx))),
+                           done_);
+
+      se::DeviceMemoryBase dst_ptr(
+          const_cast<char*>(buffer_input->tensor_data().data()) + offset_bytes,
+          input_bytes_vec->at(idx));
+      ctx->op_device_context()->stream()->ThenMemcpy(
+          &dst_ptr,
+          se::DeviceMemoryBase(
+              const_cast<char*>(n_input_list[idx].tensor_data().data()),
+              input_bytes_vec->at(idx)),
+          input_bytes_vec->at(idx));
+      offset_bytes += input_bytes_vec->at(idx);
+    }
+
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [input_bytes_vec, n_output, buffer_input, buffer_output,
+              buffer_comm_input, buffer_comm_output, ctx, coll, this, done_]() {
           auto call = functor::NcclAlltoallCall<DTYPE, WIRE_DTYPE>();
           OP_REQUIRES_OK_ASYNC(
               ctx,
-              call(*input, output, comm_input, comm_output, ctx, comm, this),
+              call(*buffer_input, buffer_output, buffer_comm_input,
+                   buffer_comm_output, topology_, ctx, coll, this),
               done_);
+          int64 offset_bytes = 0;
+          for (int idx = 0; idx < N_; ++idx) {
+            se::DeviceMemoryBase dst_ptr(
+                const_cast<char*>(n_output->at(idx)->tensor_data().data()),
+                input_bytes_vec->at(idx));
+            coll->stream()->ThenMemcpy(
+                &dst_ptr,
+                se::DeviceMemoryBase(
+                    const_cast<char*>(buffer_output->tensor_data().data()) +
+                        offset_bytes,
+                    input_bytes_vec->at(idx)),
+                input_bytes_vec->at(idx));
+            offset_bytes += input_bytes_vec->at(idx);
+          }
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
         });
   }
+
+ private:
+  CollectiveTopology topology_;
+  int64 N_;
 };
 
 #define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
-  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoall")                         \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoallMergedN")                  \
                               .Device(DEVICE_GPU)                        \
                               .TypeConstraint<DTYPE>("dtype")            \
                               .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
-                          NcclAlltoallOp<DTYPE, WIRE_DTYPE>);
+                          NcclAlltoallMergedNOp<DTYPE, WIRE_DTYPE>);
 TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 #endif  // GOOGLE_CUDA
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallv.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallv.cc
index d5338c0a..675b2deb 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallv.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallv.cc
@@ -22,13 +22,131 @@ limitations under the License.
 #include <vector>
 
 #include "hybridbackend/tensorflow/common/cast.h"
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/tensorflow/common/host_functions.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
+using GPUDevice = Eigen::GpuDevice;
+
 namespace hybridbackend {
 
 #if HYBRIDBACKEND_NCCL
 
+#if GOOGLE_CUDA
+namespace functor {
+template <typename DTYPE, typename WIRE_DTYPE>
+struct NcclAlltoallvCall {
+  Status operator()(const Tensor& input, const int32* send_counts,
+                    const int32* recv_counts, const int64 common_size,
+                    Tensor* output, Tensor* comm_input, Tensor* comm_output,
+                    CollectiveTopology topology, OpKernelContext* ctx,
+                    NcclCollective* coll, NcclCollectiveAsyncOp* comm_op) {
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name()
+            << "] [Alltoallv]";
+    TF_RETURN_IF_ERROR(coll->Alltoallv(input, send_counts, recv_counts,
+                                       common_size, output, topology));
+    return Status::OK();
+  }
+};
+
+template <>
+struct NcclAlltoallvCall<float, Eigen::half> {
+  Status operator()(const Tensor& input, const int32* send_counts,
+                    const int32* recv_counts, const int64 common_size,
+                    Tensor* output, Tensor* comm_input, Tensor* comm_output,
+                    CollectiveTopology topology, OpKernelContext* ctx,
+                    NcclCollective* coll, NcclCollectiveAsyncOp* comm_op) {
+    TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, input.shape(), comm_input));
+    TF_RETURN_IF_ERROR(
+        ctx->allocate_temp(DT_HALF, output->shape(), comm_output));
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastIn]";
+    functor::Cast<float, Eigen::half>()(input, comm_input, ctx,
+                                        coll->stream()->get());
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name()
+            << "] [Alltoallv]";
+    TF_RETURN_IF_ERROR(coll->Alltoallv(*comm_input, send_counts, recv_counts,
+                                       common_size, comm_output, topology));
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastOut]";
+    functor::Cast<Eigen::half, float>()(*comm_output, output, ctx,
+                                        coll->stream()->get());
+    return Status::OK();
+  }
+};
+
+template <typename DTYPE, typename WIRE_DTYPE>
+struct NcclAlltoallvNCall {
+  Status operator()(const std::vector<Tensor>& n_input,
+                    const std::vector<Tensor*>& n_send_sizes,
+                    const std::vector<Tensor*>& n_recv_sizes,
+                    const std::vector<int64>& n_common_size,
+                    std::vector<Tensor*>* n_output,
+                    std::vector<Tensor*>* n_comm_input,
+                    std::vector<Tensor*>* n_comm_output,
+                    CollectiveTopology topology, OpKernelContext* ctx,
+                    NcclCollective* coll, NcclCollectiveAsyncOp* comm_op) {
+    std::vector<int32*> n_send_sizes_vec(n_input.size(), nullptr);
+    std::vector<int32*> n_recv_sizes_vec(n_input.size(), nullptr);
+    for (size_t idx = 0; idx < n_input.size(); ++idx) {
+      n_send_sizes_vec[idx] = n_send_sizes[idx]->flat<int32>().data();
+      n_recv_sizes_vec[idx] = n_recv_sizes[idx]->flat<int32>().data();
+    }
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name()
+            << "] [AlltoallvN]";
+    TF_RETURN_IF_ERROR(coll->AlltoallvN(n_input, n_send_sizes_vec,
+                                        n_recv_sizes_vec, n_common_size,
+                                        n_output, topology));
+    return Status::OK();
+  }
+};
+
+template <>
+struct NcclAlltoallvNCall<float, Eigen::half> {
+  Status operator()(const std::vector<Tensor>& n_input,
+                    const std::vector<Tensor*>& n_send_sizes,
+                    const std::vector<Tensor*>& n_recv_sizes,
+                    const std::vector<int64>& n_common_size,
+                    std::vector<Tensor*>* n_output,
+                    std::vector<Tensor*>* n_comm_input,
+                    std::vector<Tensor*>* n_comm_output,
+                    CollectiveTopology topology, OpKernelContext* ctx,
+                    NcclCollective* coll, NcclCollectiveAsyncOp* comm_op) {
+    std::vector<int32*> n_send_sizes_vec(n_input.size(), 0);
+    std::vector<int32*> n_recv_sizes_vec(n_input.size(), 0);
+    for (size_t idx = 0; idx < n_input.size(); ++idx) {
+      n_send_sizes_vec[idx] = n_send_sizes[idx]->flat<int32>().data();
+      n_recv_sizes_vec[idx] = n_recv_sizes[idx]->flat<int32>().data();
+    }
+    for (int idx = 0; idx < n_input.size(); ++idx) {
+      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, n_input[idx].shape(),
+                                            n_comm_input->at(idx),
+                                            ctx->input_alloc_attr(idx)));
+      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, n_output->at(idx)->shape(),
+                                            n_comm_output->at(idx),
+                                            ctx->output_alloc_attr(idx)));
+    }
+    coll->stream()->ThenWaitUntilComputeDone(ctx);
+
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastNIn]";
+    functor::CastN<float, Eigen::half>()(n_input, n_comm_input, ctx,
+                                         coll->stream()->get());
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name()
+            << "] [AlltoallvN]";
+    TF_RETURN_IF_ERROR(coll->AlltoallvN(*n_comm_input, n_send_sizes_vec,
+                                        n_recv_sizes_vec, n_common_size,
+                                        n_comm_output, topology));
+    VLOG(1) << coll->DebugString() << " [" << comm_op->name() << "] [CastNOut]";
+    functor::CastN<Eigen::half, float>()(*n_comm_output, n_output, ctx,
+                                         coll->stream()->get());
+    return Status::OK();
+  }
+};
+
+}  // namespace functor
+#endif  // GOOGLE_CUDA
+
 REGISTER_OP("HbNcclAlltoallv")
     .Output("output: dtype")
     .Output("output_sizes: int32")
@@ -36,6 +154,7 @@ REGISTER_OP("HbNcclAlltoallv")
     .Input("input: dtype")
     .Input("input_sizes: int32")
     .Attr("common_shape: shape = {}")
+    .Attr("topology: int = 0")
     .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
     .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
     .SetIsStateful()
@@ -53,61 +172,24 @@ REGISTER_OP("HbNcclAlltoallv")
       return Status::OK();
     })
     .Doc(R"doc(
-AllToAllv using a NCCL communicator with merged buffer.
+AllToAllv using a NCCL communicator.
 
-output: a single merged tensor for all devices.
-output_sizes: a single tensor for 1st dim of outputs.
+output: An exchanged tensor.
+output_sizes: An tensor for 1st dim of the output.
 handle: Handle of a NCCL communicator.
-input: a single merged tensor for all devices.
-input_sizes: a tensor for 1st dim of inputs for all devices.
+input: A tensor to exchange.
+input_sizes: A tensor for 1st dim of input.
 )doc");
 
 #if GOOGLE_CUDA
-
-namespace functor {
 template <typename DTYPE, typename WIRE_DTYPE>
-struct NcclAlltoallvCall {
-  Status operator()(const Tensor& input, const Tensor& host_sizes,
-                    const int64 common_size, Tensor* output, Tensor* comm_input,
-                    Tensor* comm_output, OpKernelContext* ctx, NcclComm* comm,
-                    NcclCommAsyncOp* comm_op) {
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [Alltoallv]";
-    TF_RETURN_IF_ERROR(comm->Alltoallv(input, host_sizes, common_size, output));
-    return Status::OK();
-  }
-};
-
-template <>
-struct NcclAlltoallvCall<float, Eigen::half> {
-  Status operator()(const Tensor& input, const Tensor& host_sizes,
-                    const int64 common_size, Tensor* output, Tensor* comm_input,
-                    Tensor* comm_output, OpKernelContext* ctx, NcclComm* comm,
-                    NcclCommAsyncOp* comm_op) {
-    TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, input.shape(), comm_input));
-    TF_RETURN_IF_ERROR(
-        ctx->allocate_temp(DT_HALF, output->shape(), comm_output));
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastIn]";
-    functor::Cast<float, Eigen::half>()(input, comm_input, ctx, comm->stream());
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [Alltoallv]";
-    TF_RETURN_IF_ERROR(
-        comm->Alltoallv(*comm_input, host_sizes, common_size, comm_output));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastOut]";
-    functor::Cast<Eigen::half, float>()(*output, comm_output, ctx,
-                                        comm->stream());
-    return Status::OK();
-  }
-};
-
-}  // namespace functor
-
-template <typename DTYPE, typename WIRE_DTYPE>
-class NcclAlltoallvOp : public NcclCommAsyncOp {
+class NcclAlltoallvOp : public NcclCollectiveAsyncOp {
  public:
-  explicit NcclAlltoallvOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {
+  explicit NcclAlltoallvOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    int topology;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("topology", &topology));
+    topology_ = static_cast<CollectiveTopology>(topology);
     OP_REQUIRES_OK(ctx, ctx->GetAttr("common_shape", &common_shape_));
     common_size_ = 1;
     for (int64 dim = 0; dim < common_shape_.dims(); ++dim) {
@@ -115,8 +197,9 @@ class NcclAlltoallvOp : public NcclCommAsyncOp {
     }
   }
 
-  virtual void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                                    DoneCallback done) override {
+  virtual void CollectiveComputeAsync(NcclCollective* coll,
+                                      OpKernelContext* ctx,
+                                      DoneCallback done) override {
     AllocatorAttributes host_alloc_attrs;
     host_alloc_attrs.set_on_host(true);
     host_alloc_attrs.set_gpu_compatible(true);
@@ -125,67 +208,64 @@ class NcclAlltoallvOp : public NcclCommAsyncOp {
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
     const Tensor* input_sizes;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input_sizes", &input_sizes), done);
-
-    Tensor* sizes = new Tensor();
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(DT_INT32, TensorShape({comm->size() * comm->size()}),
-                           sizes, ctx->output_alloc_attr(0)),
-        done);
-
-    Tensor* host_sizes = new Tensor();
+    int active_size = coll->compute_active_size(topology_);
+    TensorShape output_sizes_shape({active_size});
+    Tensor* output_sizes;
     OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(DT_INT32, TensorShape({comm->size() * comm->size()}),
-                           host_sizes, host_alloc_attrs),
-        done);
+        ctx, ctx->allocate_output(1, output_sizes_shape, &output_sizes), done);
 
+    Tensor* host_input_sizes = new Tensor();
     Tensor* host_output_sizes = new Tensor();
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(DT_INT32, TensorShape({comm->size()}),
-                           host_output_sizes, host_alloc_attrs),
-        done);
-
     Tensor* comm_input = new Tensor();
     Tensor* comm_output = new Tensor();
-    auto done_ = [host_sizes, host_output_sizes, comm_input, comm_output,
+    auto done_ = [host_input_sizes, host_output_sizes, comm_input, comm_output,
                   done]() {
-      delete host_sizes;
+      delete host_input_sizes;
       delete host_output_sizes;
       delete comm_input;
       delete comm_output;
       done();
     };
-    comm->RunAsync(
-        "NcclAlltoallv", ctx, done_,
-        [input_sizes, sizes, host_output_sizes, input, host_sizes, comm_input,
-         comm_output, ctx, comm, this, done_]() {
-          // Collect sizes of all inputs across devices.
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Allgather]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Allgather(*input_sizes, sizes),
-                               done_);
-          ThenWaitFor(ctx, comm->ThenRecordEvent());
-          ThenCopyToHost(ctx, host_sizes, *sizes);
-          BlockHostUntilDone(ctx);
-          delete sizes;
-
-          // Allocate outputs.
-          int32 total_output_size = 0;
-          for (int32 i = 0; i < comm->size(); ++i) {
-            int32 output_size =
-                host_sizes->flat<int32>()(comm->size() * i + comm->rank());
-            total_output_size += output_size;
-            host_output_sizes->flat<int32>()(i) = output_size;
-          }
 
-          TensorShape output_sizes_shape({comm->size()});
-          Tensor* output_sizes;
+    OP_REQUIRES_OK_ASYNC(
+        ctx,
+        ctx->allocate_temp(DT_INT32, TensorShape({active_size}),
+                           host_input_sizes, host_alloc_attrs),
+        done_);
+    OP_REQUIRES_OK_ASYNC(
+        ctx,
+        ctx->allocate_temp(DT_INT32, TensorShape({active_size}),
+                           host_output_sizes, host_alloc_attrs),
+        done_);
+    coll->stream()->LaunchUntilComputeDone(
+        ctx,
+        [input_sizes, output_sizes, host_input_sizes, host_output_sizes, input,
+         comm_input, comm_output, active_size, ctx, coll, this, done_]() {
+          ctx->op_device_context()->stream()->ThenMemcpy(
+              const_cast<char*>(host_input_sizes->tensor_data().data()),
+              se::DeviceMemoryBase(
+                  const_cast<char*>(input_sizes->tensor_data().data()),
+                  input_sizes->TotalBytes()),
+              host_input_sizes->TotalBytes());
+          // Collect sizes of all inputs across devices.
+          VLOG(1) << coll->DebugString() << " [" << name() << "] [Alltoall]";
           OP_REQUIRES_OK_ASYNC(
-              ctx, ctx->allocate_output(1, output_sizes_shape, &output_sizes),
+              ctx, coll->Alltoall(*input_sizes, output_sizes, topology_),
               done_);
-          ThenCopyToDevice(ctx, output_sizes, *host_output_sizes);
+          coll->stream()->BlockComputeUntilDone(ctx);
+          ctx->op_device_context()->stream()->ThenMemcpy(
+              const_cast<char*>(host_output_sizes->tensor_data().data()),
+              se::DeviceMemoryBase(
+                  const_cast<char*>(output_sizes->tensor_data().data()),
+                  output_sizes->TotalBytes()),
+              host_output_sizes->TotalBytes());
+          ctx->op_device_context()->stream()->BlockHostUntilDone();
 
+          // Allocate output.
+          int32 total_output_size = 0;
+          for (int32 rank = 0; rank < active_size; ++rank) {
+            total_output_size += host_output_sizes->flat<int32>()(rank);
+          }
           TensorShape output_shape;
           PartialTensorShape({total_output_size})
               .Concatenate(common_shape_)
@@ -196,16 +276,20 @@ class NcclAlltoallvOp : public NcclCommAsyncOp {
 
           // Cast and communicate.
           auto call = functor::NcclAlltoallvCall<DTYPE, WIRE_DTYPE>();
-          OP_REQUIRES_OK_ASYNC(ctx,
-                               call(*input, *host_sizes, common_size_, output,
-                                    comm_input, comm_output, ctx, comm, this),
-                               done_);
+          OP_REQUIRES_OK_ASYNC(
+              ctx,
+              call(*input, host_input_sizes->flat<int32>().data(),
+                   host_output_sizes->flat<int32>().data(), common_size_,
+                   output, comm_input, comm_output, topology_, ctx, coll, this),
+              done_);
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
         });
   }
 
  private:
   PartialTensorShape common_shape_;
   int64 common_size_;
+  CollectiveTopology topology_;
 };
 
 #define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
@@ -218,7 +302,478 @@ TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 #endif  // GOOGLE_CUDA
 
-#endif
+REGISTER_OP("HbNcclAlltoallvN")
+    .Output("n_output: N * dtype")
+    .Output("n_output_sizes: N * int32")
+    .Input("handle: resource")
+    .Input("n_input: N * dtype")
+    .Input("n_input_sizes: N * int32")
+    .Attr("N: int >= 1 = 1")
+    .Attr("common_shape: list(shape)")
+    .Attr("topology: int = 0")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
+    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      int64 N;
+      TF_RETURN_IF_ERROR(c->GetAttr("N", &N));
+      std::vector<PartialTensorShape> common_shape;
+      TF_RETURN_IF_ERROR(c->GetAttr("common_shape", &common_shape));
+      for (int64 n = 0; n < N; ++n) {
+        shape_inference::ShapeHandle shape;
+        TF_RETURN_IF_ERROR(
+            c->MakeShapeFromPartialTensorShape(common_shape[n], &shape));
+        TF_RETURN_IF_ERROR(c->Concatenate(
+            c->Vector(shape_inference::InferenceContext::kUnknownDim), shape,
+            &shape));
+        c->set_output(n, shape);
+        c->set_output(N + n, c->input(1 + N + n));
+      }
+      return Status::OK();
+    })
+    .Doc(R"doc(
+Packed AllToAllv using a NCCL communicator.
+
+n_output: N exchanged tensors.
+n_output_sizes: N tensors for 1st dim of outputs.
+handle: Handle of a NCCL communicator.
+n_input: N tensors to exchange.
+n_input_sizes: N tensors for 1st dim of inputs.
+)doc");
+
+#if GOOGLE_CUDA
+template <typename DTYPE, typename WIRE_DTYPE>
+class NcclAlltoallvNOp : public NcclCollectiveAsyncOp {
+ public:
+  explicit NcclAlltoallvNOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &N_));
+    int topology;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("topology", &topology));
+    topology_ = static_cast<CollectiveTopology>(topology);
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("common_shape", &common_shape_));
+    for (int64 idx = 0; idx < N_; ++idx) {
+      int64 common_shape_size = 1;
+      for (int64 dim = 0; dim < common_shape_[idx].dims(); ++dim) {
+        common_shape_size *= common_shape_[idx].dim_size(dim);
+      }
+      common_sizes_.push_back(common_shape_size);
+    }
+  }
+
+  virtual void CollectiveComputeAsync(NcclCollective* coll,
+                                      OpKernelContext* ctx,
+                                      DoneCallback done) override {
+    AllocatorAttributes host_alloc_attrs;
+    host_alloc_attrs.set_on_host(true);
+    host_alloc_attrs.set_gpu_compatible(true);
+
+    std::vector<Tensor>* n_input = new std::vector<Tensor>();
+    std::vector<Tensor>* n_input_sizes = new std::vector<Tensor>();
+    std::vector<Tensor*>* n_output_sizes = new std::vector<Tensor*>();
+    std::vector<Tensor*>* n_host_input_sizes = new std::vector<Tensor*>();
+    std::vector<Tensor*>* n_host_output_sizes = new std::vector<Tensor*>();
+    std::vector<Tensor*>* n_comm_input = new std::vector<Tensor*>();
+    std::vector<Tensor*>* n_comm_output = new std::vector<Tensor*>();
+    for (int idx = 0; idx < N_; ++idx) {
+      n_host_input_sizes->push_back(new Tensor());
+      n_host_output_sizes->push_back(new Tensor());
+      n_comm_input->push_back(new Tensor());
+      n_comm_output->push_back(new Tensor());
+    }
+    auto done_ = [this, n_input, n_input_sizes, n_output_sizes,
+                  n_host_input_sizes, n_host_output_sizes, n_comm_input,
+                  n_comm_output, done]() {
+      delete n_input;
+      delete n_input_sizes;
+      delete n_output_sizes;
+      for (int idx = 0; idx < N_; ++idx) {
+        delete n_host_input_sizes->at(idx);
+        delete n_host_output_sizes->at(idx);
+        delete n_comm_input->at(idx);
+        delete n_comm_output->at(idx);
+      }
+      delete n_host_input_sizes;
+      delete n_host_output_sizes;
+      delete n_comm_input;
+      delete n_comm_output;
+      done();
+    };
+
+    OpInputList n_input_list;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("n_input", &n_input_list), done_);
+    OpInputList n_input_sizes_list;
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->input_list("n_input_sizes", &n_input_sizes_list), done_);
+
+    int active_size = coll->compute_active_size(topology_);
+
+    for (int idx = 0; idx < N_; ++idx) {
+      n_input->push_back(n_input_list[idx]);
+      auto& input_sizes = n_input_sizes_list[idx];
+      OP_REQUIRES_ASYNC(
+          ctx, input_sizes.NumElements() == active_size,
+          errors::InvalidArgument(
+              "Sizes of input ", idx, " has ", input_sizes.NumElements(),
+              " elements, which is not equal to active rank size: ",
+              active_size),
+          done_);
+      n_input_sizes->push_back(input_sizes);
+      OP_REQUIRES_OK_ASYNC(
+          ctx,
+          ctx->allocate_temp(DT_INT32, TensorShape({active_size}),
+                             n_host_input_sizes->at(idx), host_alloc_attrs),
+          done);
+      OP_REQUIRES_OK_ASYNC(
+          ctx,
+          ctx->allocate_temp(DT_INT32, TensorShape({active_size}),
+                             n_host_output_sizes->at(idx), host_alloc_attrs),
+          done);
+      Tensor* output_sizes;
+      OP_REQUIRES_OK_ASYNC(
+          ctx, ctx->allocate_output(N_ + idx, {active_size}, &output_sizes),
+          done_);
+      n_output_sizes->push_back(output_sizes);
+    }
+
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [this, coll, ctx, done_, n_input, n_input_sizes, n_output_sizes,
+              n_host_input_sizes, n_host_output_sizes, n_comm_input,
+              n_comm_output, active_size]() {
+          for (size_t idx = 0; idx < N_; ++idx) {
+            ctx->op_device_context()->stream()->ThenMemcpy(
+                const_cast<char*>(
+                    n_host_input_sizes->at(idx)->tensor_data().data()),
+                se::DeviceMemoryBase(
+                    const_cast<char*>(
+                        n_input_sizes->at(idx).tensor_data().data()),
+                    n_input_sizes->at(idx).TotalBytes()),
+                n_host_input_sizes->at(idx)->TotalBytes());
+          }
+          VLOG(1) << coll->DebugString() << " [" << name() << "] [AlltoallN]";
+          OP_REQUIRES_OK_ASYNC(
+              ctx, coll->AlltoallN(*n_input_sizes, n_output_sizes, topology_),
+              done_);
+          coll->stream()->BlockComputeUntilDone(ctx);
+          for (size_t idx = 0; idx < N_; ++idx) {
+            ctx->op_device_context()->stream()->ThenMemcpy(
+                const_cast<char*>(
+                    n_host_output_sizes->at(idx)->tensor_data().data()),
+                se::DeviceMemoryBase(
+                    const_cast<char*>(
+                        n_output_sizes->at(idx)->tensor_data().data()),
+                    n_output_sizes->at(idx)->TotalBytes()),
+                n_host_output_sizes->at(idx)->TotalBytes());
+          }
+          ctx->op_device_context()->stream()->BlockHostUntilDone();
+          std::vector<int32> n_host_output_total_size(N_, 0);
+          for (size_t idx = 0; idx < N_; ++idx) {
+            n_host_output_total_size[idx] = 0;
+            for (int32 rank = 0; rank < active_size; ++rank) {
+              n_host_output_total_size[idx] +=
+                  n_host_output_sizes->at(idx)->flat<int32>()(rank);
+            }
+          }
+
+          std::vector<Tensor*> n_output;
+          for (int idx = 0; idx < N_; ++idx) {
+            TensorShape output_shape;
+            PartialTensorShape({n_host_output_total_size[idx]})
+                .Concatenate(common_shape_[idx])
+                .AsTensorShape(&output_shape);
+            Tensor* output;
+            OP_REQUIRES_OK_ASYNC(
+                ctx, ctx->allocate_output(idx, output_shape, &output), done_);
+            n_output.push_back(output);
+          }
+
+          auto call = functor::NcclAlltoallvNCall<DTYPE, WIRE_DTYPE>();
+          OP_REQUIRES_OK_ASYNC(
+              ctx,
+              call(*n_input, *n_host_input_sizes, *n_host_output_sizes,
+                   common_sizes_, &n_output, n_comm_input, n_comm_output,
+                   topology_, ctx, coll, this),
+              done_);
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
+        });
+  }
+
+ private:
+  CollectiveTopology topology_;
+  int64 N_;
+  std::vector<PartialTensorShape> common_shape_;
+  std::vector<int64> common_sizes_;
+};
+
+#define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoallvN")                       \
+                              .Device(DEVICE_GPU)                        \
+                              .TypeConstraint<DTYPE>("dtype")            \
+                              .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
+                          NcclAlltoallvNOp<DTYPE, WIRE_DTYPE>);
+TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
+#undef REGISTER_KERNEL
+#endif  // GOOGLE_CUDA
+
+REGISTER_OP("HbNcclAlltoallvMergedN")
+    .Output("n_output: N * dtype")
+    .Output("n_output_sizes: N * int32")
+    .Input("handle: resource")
+    .Input("n_input: N * dtype")
+    .Input("n_input_sizes: N * int32")
+    .Attr("N: int >= 1 = 1")
+    .Attr("common_shape: shape = {}")
+    .Attr("topology: int = 0")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
+    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      int64 N;
+      TF_RETURN_IF_ERROR(c->GetAttr("N", &N));
+
+      PartialTensorShape common_shape;
+      TF_RETURN_IF_ERROR(c->GetAttr("common_shape", &common_shape));
+      for (int64 n = 0; n < N; ++n) {
+        shape_inference::ShapeHandle shape;
+        TF_RETURN_IF_ERROR(
+            c->MakeShapeFromPartialTensorShape(common_shape, &shape));
+        TF_RETURN_IF_ERROR(c->Concatenate(
+            c->Vector(shape_inference::InferenceContext::kUnknownDim), shape,
+            &shape));
+        c->set_output(n, shape);
+        c->set_output(N + n, c->input(1 + N + n));
+      }
+      return Status::OK();
+    })
+    .Doc(R"doc(
+Packed merged AllToAllv using a NCCL communicator.
+
+n_output: N exchanged tensors.
+n_output_sizes: N tensors for 1st dim of outputs.
+handle: Handle of a NCCL communicator.
+n_input: N tensors to exchange.
+n_input_sizes: N tensors for 1st dim of inputs.
+)doc");
+
+#if GOOGLE_CUDA
+template <typename DTYPE, typename WIRE_DTYPE>
+class NcclAlltoallvMergedNOp : public NcclCollectiveAsyncOp {
+ public:
+  explicit NcclAlltoallvMergedNOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &N_));
+    int topology;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("topology", &topology));
+    topology_ = static_cast<CollectiveTopology>(topology);
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("common_shape", &common_shape_));
+    common_size_ = 1;
+    for (int64 dim = 0; dim < common_shape_.dims(); ++dim) {
+      common_size_ *= common_shape_.dim_size(dim);
+    }
+  }
+
+  virtual void CollectiveComputeAsync(NcclCollective* coll,
+                                      OpKernelContext* ctx,
+                                      DoneCallback done) override {
+    // FIXME: Random crash happens
+    AllocatorAttributes host_alloc_attrs;
+    host_alloc_attrs.set_on_host(true);
+    host_alloc_attrs.set_gpu_compatible(true);
+
+    std::vector<Tensor>* n_input_sizes = new std::vector<Tensor>();
+    std::vector<Tensor*>* n_output_sizes = new std::vector<Tensor*>();
+    Tensor* buffer_input = new Tensor();
+    Tensor* buffer_output = new Tensor();
+    Tensor* buffer_comm_input = new Tensor();
+    Tensor* buffer_comm_output = new Tensor();
+    Tensor* host_sizes_buffer = new Tensor();
+    auto done_ = [n_input_sizes, n_output_sizes, buffer_input, buffer_output,
+                  buffer_comm_input, buffer_comm_output, host_sizes_buffer,
+                  done]() {
+      delete n_input_sizes;
+      delete n_output_sizes;
+      delete buffer_input;
+      delete buffer_output;
+      delete buffer_comm_input;
+      delete buffer_comm_output;
+      delete host_sizes_buffer;
+      done();
+    };
+
+    OpInputList n_input_list;
+    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("n_input", &n_input_list), done_);
+    OpInputList n_input_sizes_list;
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->input_list("n_input_sizes", &n_input_sizes_list), done_);
+    std::vector<int64> input_bytes_vec(N_, 0);
+    int64 total_input_bytes = 0;
+    int active_size = coll->compute_active_size(topology_);
+    for (int idx = 0; idx < N_; ++idx) {
+      auto& input_sizes = n_input_sizes_list[idx];
+      OP_REQUIRES_ASYNC(
+          ctx, input_sizes.NumElements() == active_size,
+          errors::InvalidArgument(
+              "Sizes of input ", idx, " has ", input_sizes.NumElements(),
+              " elements, which is not equal to active rank size: ",
+              active_size),
+          done_);
+      n_input_sizes->push_back(input_sizes);
+      Tensor* output_sizes;
+      OP_REQUIRES_OK_ASYNC(
+          ctx, ctx->allocate_output(N_ + idx, {active_size}, &output_sizes),
+          done_);
+      n_output_sizes->push_back(output_sizes);
+      const auto input_bytes = n_input_list[idx].TotalBytes();
+      total_input_bytes += input_bytes;
+      input_bytes_vec[idx] = input_bytes;
+    }
+    OP_REQUIRES_OK_ASYNC(ctx,
+                         ctx->allocate_temp(DataTypeToEnum<DTYPE>::value,
+                                            {total_input_bytes}, buffer_input),
+                         done_);
+    int64 input_offset_bytes = 0;
+    for (int64 idx = 0; idx < N_; ++idx) {
+      se::DeviceMemoryBase input_dst_ptr(
+          const_cast<char*>(buffer_input->tensor_data().data()) +
+              input_offset_bytes,
+          input_bytes_vec[idx]);
+      ctx->op_device_context()->stream()->ThenMemcpy(
+          &input_dst_ptr,
+          se::DeviceMemoryBase(
+              const_cast<char*>(n_input_list[idx].tensor_data().data()),
+              input_bytes_vec[idx]),
+          input_bytes_vec[idx]);
+      input_offset_bytes += input_bytes_vec[idx];
+    }
+
+    OP_REQUIRES_OK_ASYNC(
+        ctx,
+        ctx->allocate_temp(DT_INT32, {(2 * N_ + 2) * active_size + N_},
+                           host_sizes_buffer, host_alloc_attrs),
+        done_);
+    int32* host_input_sizes = host_sizes_buffer->flat<int32>().data();
+    int32* host_output_sizes =
+        host_sizes_buffer->flat<int32>().data() + N_ * active_size;
+    int32* host_input_total_sizes =
+        host_sizes_buffer->flat<int32>().data() + 2 * N_ * active_size;
+    int32* host_output_total_sizes =
+        host_sizes_buffer->flat<int32>().data() + (2 * N_ + 1) * active_size;
+    int32* host_n_output_sizes =
+        host_sizes_buffer->flat<int32>().data() + (2 * N_ + 2) * active_size;
+
+    coll->stream()->LaunchUntilComputeDone(
+        ctx,
+        [n_input_sizes, n_output_sizes, host_input_sizes, host_output_sizes,
+         host_input_total_sizes, host_output_total_sizes, host_n_output_sizes,
+         buffer_input, buffer_output, buffer_comm_input, buffer_comm_output,
+         active_size, coll, ctx, done_, this]() {
+          VLOG(1) << coll->DebugString() << " [" << name() << "] [AlltoallN]";
+          OP_REQUIRES_OK_ASYNC(
+              ctx, coll->AlltoallN(*n_input_sizes, n_output_sizes, topology_),
+              done_);
+          coll->stream()->BlockComputeUntilDone(ctx);
+          for (int64 idx = 0; idx < N_; ++idx) {
+            ctx->op_device_context()->stream()->ThenMemcpy(
+                host_input_sizes + idx * active_size,
+                se::DeviceMemoryBase(
+                    n_input_sizes->at(idx).flat<int32>().data(),
+                    n_input_sizes->at(idx).TotalBytes()),
+                n_input_sizes->at(idx).TotalBytes());
+            ctx->op_device_context()->stream()->ThenMemcpy(
+                host_output_sizes + idx * active_size,
+                se::DeviceMemoryBase(
+                    n_output_sizes->at(idx)->flat<int32>().data(),
+                    n_output_sizes->at(idx)->TotalBytes()),
+                n_output_sizes->at(idx)->TotalBytes());
+          }
+          ctx->op_device_context()->stream()->BlockHostUntilDone();
+          for (int64 rank = 0; rank < active_size; ++rank) {
+            host_input_total_sizes[rank] = 0;
+            host_output_total_sizes[rank] = 0;
+            for (int64 idx = 0; idx < N_; ++idx) {
+              host_input_total_sizes[rank] +=
+                  host_input_sizes[idx * active_size + rank];
+              host_output_total_sizes[rank] +=
+                  host_output_sizes[idx * active_size + rank];
+            }
+          }
+          for (int64 idx = 0; idx < N_; ++idx) {
+            host_n_output_sizes[idx] = 0;
+            for (int64 rank = 0; rank < active_size; ++rank) {
+              host_n_output_sizes[idx] +=
+                  host_output_sizes[idx * active_size + rank];
+            }
+          }
+
+          int64 total_output_sizes = 0;
+          for (int rank = 0; rank < active_size; ++rank) {
+            total_output_sizes += host_output_total_sizes[rank];
+          }
+          OP_REQUIRES_OK_ASYNC(
+              ctx,
+              ctx->allocate_temp(DataTypeToEnum<DTYPE>::value,
+                                 {common_size_ * total_output_sizes},
+                                 buffer_output),
+              done_);
+
+          auto call = functor::NcclAlltoallvCall<DTYPE, WIRE_DTYPE>();
+          OP_REQUIRES_OK_ASYNC(
+              ctx,
+              call(*buffer_input, host_input_total_sizes,
+                   host_output_total_sizes, common_size_, buffer_output,
+                   buffer_comm_input, buffer_comm_output, topology_, ctx, coll,
+                   this),
+              done_);
+          coll->stream()->BlockComputeUntilDone(ctx);
+
+          std::vector<Tensor*> n_output;
+          for (int idx = 0; idx < N_; ++idx) {
+            TensorShape output_shape;
+            PartialTensorShape({host_n_output_sizes[idx]})
+                .Concatenate(common_shape_)
+                .AsTensorShape(&output_shape);
+            Tensor* output;
+            OP_REQUIRES_OK_ASYNC(
+                ctx, ctx->allocate_output(idx, output_shape, &output), done_);
+            n_output.push_back(output);
+          }
+          int64 output_offset_bytes = 0;
+          for (int idx = 0; idx < N_; ++idx) {
+            se::DeviceMemoryBase output_dst_ptr(
+                const_cast<char*>(n_output[idx]->tensor_data().data()),
+                host_n_output_sizes[idx]);
+            ctx->op_device_context()->stream()->ThenMemcpy(
+                &output_dst_ptr,
+                se::DeviceMemoryBase(
+                    const_cast<char*>(buffer_output->tensor_data().data()) +
+                        output_offset_bytes,
+                    host_n_output_sizes[idx]),
+                host_n_output_sizes[idx]);
+            output_offset_bytes += host_n_output_sizes[idx] * common_size_ *
+                                   DataTypeSize(DataTypeToEnum<DTYPE>::value);
+          }
+          coll->stream()->BlockComputeUntilDone(ctx, done_);
+        });
+  }
+
+ private:
+  CollectiveTopology topology_;
+  int64 N_;
+  PartialTensorShape common_shape_;
+  int64 common_size_;
+};
+
+#define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoallvMergedN")                 \
+                              .Device(DEVICE_GPU)                        \
+                              .TypeConstraint<DTYPE>("dtype")            \
+                              .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
+                          NcclAlltoallvMergedNOp<DTYPE, WIRE_DTYPE>);
+TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
+#undef REGISTER_KERNEL
+#endif  // GOOGLE_CUDA
+
+#endif  // HYBRIDBACKEND_NCCL
 
 }  // namespace hybridbackend
 }  // namespace tensorflow
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallv_n.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallv_n.cc
deleted file mode 100644
index 1cd5b6ca..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallv_n.cc
+++ /dev/null
@@ -1,335 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#include <tensorflow/core/framework/common_shape_fns.h>
-#include <tensorflow/core/framework/op_kernel.h>
-#include <tensorflow/core/framework/shape_inference.h>
-
-#include <vector>
-
-#include "hybridbackend/tensorflow/common/cast.h"
-#include "hybridbackend/tensorflow/common/host_functions.h"
-#include "hybridbackend/tensorflow/common/slice_sum.h"
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
-
-namespace tensorflow {
-using GPUDevice = Eigen::GpuDevice;
-
-namespace hybridbackend {
-
-#if HYBRIDBACKEND_NCCL
-
-REGISTER_OP("HbNcclAlltoallvN")
-    .Output("outputs: num_columns * dtype")
-    .Output("outputs_sizes: num_columns * int32")
-    .Input("handle: resource")
-    .Input("inputs: num_columns * dtype")
-    .Input("inputs_sizes: num_columns * int32")
-    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
-    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
-    .Attr("num_columns: int >= 1 = 1")
-    .Attr("common_shapes: list(shape)")
-    .SetIsStateful()
-    .SetShapeFn([](shape_inference::InferenceContext* c) {
-      int64 num_columns;
-      TF_RETURN_IF_ERROR(c->GetAttr("num_columns", &num_columns));
-      std::vector<PartialTensorShape> common_shapes;
-      TF_RETURN_IF_ERROR(c->GetAttr("common_shapes", &common_shapes));
-      for (int64 n = 0; n < num_columns; ++n) {
-        shape_inference::ShapeHandle shape;
-        TF_RETURN_IF_ERROR(
-            c->MakeShapeFromPartialTensorShape(common_shapes[n], &shape));
-        TF_RETURN_IF_ERROR(c->Concatenate(
-            c->Vector(shape_inference::InferenceContext::kUnknownDim), shape,
-            &shape));
-        c->set_output(n, shape);
-        c->set_output(num_columns + n, c->input(1 + num_columns + n));
-      }
-      return Status::OK();
-    })
-    .Doc(R"doc(
-Grouped AllToAllv using a NCCL communicator with merged buffer.
-
-outputs: Merged tensors for all devices.
-outputs_sizes: Tensors for 1st dim of outputs.
-handle: Handle of a NCCL communicator.
-inputs: Merged tensors for all devices.
-inputs_sizes: Tensors for 1st dim of inputs for all devices.
-)doc");
-
-#if GOOGLE_CUDA
-
-namespace functor {
-template <typename DTYPE, typename WIRE_DTYPE>
-struct NcclAlltoallvNCall {
-  Status operator()(const std::vector<Tensor>& inputs, const Tensor& host_sizes,
-                    const std::vector<int64>& common_sizes,
-                    std::vector<Tensor*>* outputs,
-                    std::vector<Tensor*>* comm_in_vec,
-                    std::vector<Tensor*>* comm_out_vec, OpKernelContext* ctx,
-                    NcclComm* comm, NcclCommAsyncOp* comm_op) {
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [AlltoallvN]";
-    TF_RETURN_IF_ERROR(
-        comm->AlltoallvN(inputs, host_sizes, common_sizes, outputs));
-    return Status::OK();
-  }
-};
-
-template <>
-struct NcclAlltoallvNCall<float, Eigen::half> {
-  Status operator()(const std::vector<Tensor>& inputs, const Tensor& host_sizes,
-                    const std::vector<int64>& common_sizes,
-                    std::vector<Tensor*>* outputs,
-                    std::vector<Tensor*>* comm_input_vec,
-                    std::vector<Tensor*>* comm_output_vec, OpKernelContext* ctx,
-                    NcclComm* comm, NcclCommAsyncOp* comm_op) {
-    for (int i = 0; i < inputs.size(); ++i) {
-      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, inputs[i].shape(),
-                                            comm_input_vec->at(i),
-                                            ctx->input_alloc_attr(i)));
-      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, outputs->at(i)->shape(),
-                                            comm_output_vec->at(i),
-                                            ctx->output_alloc_attr(i)));
-    }
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastNIn]";
-    functor::CastN<float, Eigen::half>()(inputs, comm_input_vec, ctx,
-                                         comm->stream());
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [AlltoallvN]";
-    TF_RETURN_IF_ERROR(comm->AlltoallvN(*comm_input_vec, host_sizes,
-                                        common_sizes, comm_output_vec));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastNOut]";
-    functor::CastN<Eigen::half, float>()(*comm_output_vec, outputs, ctx,
-                                         comm->stream());
-    return Status::OK();
-  }
-};
-
-}  // namespace functor
-
-template <typename DTYPE, typename WIRE_DTYPE>
-class NcclAlltoallvNOp : public NcclCommAsyncOp {
- public:
-  explicit NcclAlltoallvNOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_columns", &num_columns_));
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("common_shapes", &common_shapes_));
-    for (int64 i = 0; i < num_columns_; ++i) {
-      int64 common_shape_size = 1;
-      for (int64 dim = 0; dim < common_shapes_[i].dims(); ++dim) {
-        common_shape_size *= common_shapes_[i].dim_size(dim);
-      }
-      common_sizes_.push_back(common_shape_size);
-    }
-  }
-
-  virtual void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                                    DoneCallback done) override {
-    AllocatorAttributes host_alloc_attrs;
-    host_alloc_attrs.set_on_host(true);
-    host_alloc_attrs.set_gpu_compatible(true);
-
-    std::vector<Tensor>* inputs_vec = new std::vector<Tensor>();
-    std::vector<Tensor>* inputs_sizes_vec = new std::vector<Tensor>();
-    std::vector<Tensor*>* outputs_sizes_vec = new std::vector<Tensor*>();
-    Tensor* all_sizes = new Tensor();
-    Tensor* host_all_sizes = new Tensor();
-    Tensor* outputs_total_size = new Tensor();
-    Tensor* host_outputs_total_size = new Tensor();
-    Tensor* host_outputs_sizes_ptrs = new Tensor();
-    std::vector<Tensor*>* comm_input_vec = new std::vector<Tensor*>();
-    for (int i = 0; i < num_columns_; ++i) {
-      comm_input_vec->push_back(new Tensor());
-    }
-    std::vector<Tensor*>* comm_output_vec = new std::vector<Tensor*>();
-    for (int i = 0; i < num_columns_; ++i) {
-      comm_output_vec->push_back(new Tensor());
-    }
-    auto done_ = [this, inputs_vec, inputs_sizes_vec, outputs_sizes_vec,
-                  all_sizes, host_all_sizes, outputs_total_size,
-                  host_outputs_total_size, host_outputs_sizes_ptrs,
-                  comm_input_vec, comm_output_vec, done]() {
-      delete inputs_vec;
-      delete inputs_sizes_vec;
-      delete outputs_sizes_vec;
-      delete all_sizes;
-      delete host_all_sizes;
-      delete outputs_total_size;
-      delete host_outputs_total_size;
-      delete host_outputs_sizes_ptrs;
-      for (int i = 0; i < num_columns_; ++i) {
-        delete comm_input_vec->at(i);
-      }
-      delete comm_input_vec;
-      for (int i = 0; i < num_columns_; ++i) {
-        delete comm_output_vec->at(i);
-      }
-      delete comm_output_vec;
-      done();
-    };
-
-    auto* ctx_stream = ctx->op_device_context()->stream();
-    CudaStream ctx_cu_stream = CudaStream(ctx_stream);
-
-    OpInputList inputs;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("inputs", &inputs), done_);
-    for (int i = 0; i < num_columns_; ++i) {
-      inputs_vec->push_back(inputs[i]);
-    }
-
-    OpInputList inputs_sizes;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("inputs_sizes", &inputs_sizes),
-                         done_);
-    for (int i = 0; i < num_columns_; ++i) {
-      auto& input_sizes = inputs_sizes[i];
-      OP_REQUIRES_ASYNC(
-          ctx, input_sizes.NumElements() == comm->size(),
-          errors::InvalidArgument(
-              "Sizes of input ", i, " has ", input_sizes.NumElements(),
-              " elements, which is not equal to communicator size: ",
-              comm->size()),
-          done_);
-      inputs_sizes_vec->push_back(input_sizes);
-    }
-
-    for (int i = 0; i < num_columns_; ++i) {
-      Tensor* outputs_sizes;
-      OP_REQUIRES_OK_ASYNC(ctx,
-                           ctx->allocate_output(num_columns_ + i,
-                                                {comm->size()}, &outputs_sizes),
-                           done_);
-      outputs_sizes_vec->push_back(outputs_sizes);
-    }
-
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(
-            DT_INT32, TensorShape({num_columns_ * comm->size() * comm->size()}),
-            all_sizes),
-        done_);
-
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(
-            DT_INT32, TensorShape({num_columns_ * comm->size() * comm->size()}),
-            host_all_sizes, host_alloc_attrs),
-        done_);
-
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(DT_INT32, TensorShape({num_columns_}),
-                           outputs_total_size),
-        done_);
-
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(DT_INT32, TensorShape({num_columns_}),
-                           host_outputs_total_size, host_alloc_attrs),
-        done_);
-
-    OP_REQUIRES_OK_ASYNC(
-        ctx,
-        ctx->allocate_temp(
-            DT_INT8,
-            TensorShape({num_columns_ * static_cast<int32>(sizeof(int32*))}),
-            host_outputs_sizes_ptrs, host_alloc_attrs),
-        done_);
-
-    comm->RunAsync(
-        "NcclAlltoallvN", ctx, done_,
-        [this, comm, ctx, ctx_cu_stream, done_, inputs_vec, inputs_sizes_vec,
-         outputs_sizes_vec, all_sizes, host_all_sizes, outputs_total_size,
-         host_outputs_total_size, host_outputs_sizes_ptrs, comm_input_vec,
-         comm_output_vec]() {
-          // Collect sizes of all inputs across devices.
-          VLOG(1) << comm->DebugString() << " [" << name()
-                  << "] [GroupAllgather]";
-          OP_REQUIRES_OK_ASYNC(
-              ctx, comm->GroupAllgather(*inputs_sizes_vec, all_sizes), done_);
-          ThenWaitFor(ctx, comm->ThenRecordEvent());
-
-          // Compute output sizes.
-          const int32* d_all_sizes = all_sizes->flat<int32>().data();
-          int32* d_outputs_total_size =
-              const_cast<int32*>(outputs_total_size->flat<int32>().data());
-          int32* h_outputs_total_size =
-              const_cast<int32*>(host_outputs_total_size->flat<int32>().data());
-          int32** hd_outputs_sizes = reinterpret_cast<int32**>(
-              host_outputs_sizes_ptrs->flat<int8>().data());
-          for (size_t i = 0; i < num_columns_; ++i) {
-            hd_outputs_sizes[i] =
-                outputs_sizes_vec->at(i)->flat<int32>().data();
-          }
-          ctx_cu_stream.ThenMemset(d_outputs_total_size, 0,
-                                   num_columns_ * sizeof(int32));
-          functor::SliceSumN<GPUDevice, int32>()(
-              comm->size(), comm->size(), comm->rank(),
-              outputs_sizes_vec->size(), d_all_sizes, d_outputs_total_size,
-              hd_outputs_sizes, ctx->eigen_device<GPUDevice>());
-          ThenCopyToHost(ctx, host_outputs_total_size, *outputs_total_size);
-          ThenCopyToHost(ctx, host_all_sizes, *all_sizes);
-          BlockHostUntilDone(ctx);
-
-          // Allocate outputs.
-          std::vector<Tensor*> outputs_vec;
-          for (int n = 0; n < num_columns_; ++n) {
-            TensorShape output_shape;
-            PartialTensorShape({h_outputs_total_size[n]})
-                .Concatenate(common_shapes_[n])
-                .AsTensorShape(&output_shape);
-            Tensor* output;
-            OP_REQUIRES_OK_ASYNC(
-                ctx, ctx->allocate_output(n, output_shape, &output), done_);
-            outputs_vec.push_back(output);
-          }
-
-          // Cast and communicate.
-          auto call = functor::NcclAlltoallvNCall<DTYPE, WIRE_DTYPE>();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              call(*inputs_vec, *host_all_sizes, common_sizes_, &outputs_vec,
-                   comm_input_vec, comm_output_vec, ctx, comm, this),
-              done_);
-        });
-  }
-
- private:
-  int64 num_columns_;
-  std::vector<PartialTensorShape> common_shapes_;
-  std::vector<int64> common_sizes_;
-};  // namespace functor
-
-#define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
-  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoallvN")                       \
-                              .Device(DEVICE_GPU)                        \
-                              .TypeConstraint<DTYPE>("dtype")            \
-                              .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
-                          NcclAlltoallvNOp<DTYPE, WIRE_DTYPE>);
-TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
-#undef REGISTER_KERNEL
-
-#endif  // GOOGLE_CUDA
-
-#endif
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallw.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallw.cc
deleted file mode 100644
index e987e789..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallw.cc
+++ /dev/null
@@ -1,280 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#include <tensorflow/core/framework/common_shape_fns.h>
-#include <tensorflow/core/framework/op_kernel.h>
-#include <tensorflow/core/framework/shape_inference.h>
-
-#include <vector>
-
-#include "hybridbackend/tensorflow/common/cast.h"
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
-
-namespace tensorflow {
-namespace hybridbackend {
-
-#if HYBRIDBACKEND_NCCL
-
-REGISTER_OP("HbNcclAlltoallw")
-    .Output("outputs: num_shards * dtype")
-    .Input("handle: resource")
-    .Input("inputs: num_shards * dtype")
-    .Attr("common_shape: shape = {}")
-    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
-    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
-    .Attr("num_shards: int >= 1 = 1")
-    .SetIsStateful()
-    .SetShapeFn([](shape_inference::InferenceContext* c) {
-      PartialTensorShape common_shape;
-      TF_RETURN_IF_ERROR(c->GetAttr("common_shape", &common_shape));
-      shape_inference::ShapeHandle shape;
-      TF_RETURN_IF_ERROR(
-          c->MakeShapeFromPartialTensorShape(common_shape, &shape));
-      TF_RETURN_IF_ERROR(c->Concatenate(
-          c->Vector(shape_inference::InferenceContext::kUnknownDim), shape,
-          &shape));
-      for (int64 dim = 0; dim < c->num_outputs(); ++dim) {
-        c->set_output(dim, shape);
-      }
-      return Status::OK();
-    })
-    .Doc(R"doc(
-AllToAllw using a NCCL communicator.
-
-outputs: Rotated tensors for each device.
-handle: Handle of a NCCL communicator.
-inputs: Tensors to rotate for each device.
-)doc");
-
-#if GOOGLE_CUDA
-
-namespace functor {
-template <typename DTYPE, typename WIRE_DTYPE>
-struct NcclAlltoallwCall {
-  Status operator()(const std::vector<Tensor>& inputs_vec,
-                    std::vector<Tensor*>* outputs_vec,
-                    std::vector<Tensor*>* comm_input_vec,
-                    std::vector<Tensor*>* comm_output_vec, OpKernelContext* ctx,
-                    NcclComm* comm, NcclCommAsyncOp* comm_op) {
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [Alltoallw]";
-    TF_RETURN_IF_ERROR(comm->Alltoallw(inputs_vec, outputs_vec));
-    return Status::OK();
-  }
-};
-
-template <>
-struct NcclAlltoallwCall<float, Eigen::half> {
-  Status operator()(const std::vector<Tensor>& inputs_vec,
-                    std::vector<Tensor*>* outputs_vec,
-                    std::vector<Tensor*>* comm_input_vec,
-                    std::vector<Tensor*>* comm_output_vec, OpKernelContext* ctx,
-                    NcclComm* comm, NcclCommAsyncOp* comm_op) {
-    for (int i = 0; i < comm->size(); ++i) {
-      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, inputs_vec[i].shape(),
-                                            comm_input_vec->at(i)));
-      TF_RETURN_IF_ERROR(ctx->allocate_temp(
-          DT_HALF, outputs_vec->at(i)->shape(), comm_output_vec->at(i)));
-    }
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastIn]";
-    for (int i = 0; i < inputs_vec.size(); i++) {
-      functor::Cast<float, Eigen::half>()(inputs_vec[i], comm_input_vec->at(i),
-                                          ctx, comm->stream());
-    }
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [Alltoallw]";
-    TF_RETURN_IF_ERROR(comm->Alltoallw(*comm_input_vec, comm_output_vec));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastOut]";
-    for (int i = 0; i < comm_output_vec->size(); i++) {
-      functor::Cast<Eigen::half, float>()(
-          *(comm_output_vec->at(i)), outputs_vec->at(i), ctx, comm->stream());
-    }
-    return Status::OK();
-  }
-};
-
-}  // namespace functor
-
-template <typename DTYPE, typename WIRE_DTYPE>
-class NcclAlltoallwOp : public NcclCommAsyncOp {
- public:
-  explicit NcclAlltoallwOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("common_shape", &common_shape_));
-    common_size_ = 1;
-    for (int64 dim = 0; dim < common_shape_.dims(); ++dim) {
-      common_size_ *= common_shape_.dim_size(dim);
-    }
-  }
-
-  virtual void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                                    DoneCallback done) override {
-    OpInputList inputs;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("inputs", &inputs), done);
-    std::vector<Tensor>* inputs_vec = new std::vector<Tensor>();
-    for (int i = 0; i < comm->size(); ++i) {
-      inputs_vec->push_back(inputs[i]);
-    }
-    ctx->set_output(comm->rank(), inputs[comm->rank()]);
-
-    std::vector<Tensor*>* comm_input_vec = new std::vector<Tensor*>();
-    std::vector<Tensor*>* comm_output_vec = new std::vector<Tensor*>();
-    for (int i = 0; i < comm->size(); ++i) {
-      comm_input_vec->push_back(new Tensor());
-      comm_output_vec->push_back(new Tensor());
-    }
-    auto done_ = [inputs_vec, comm_input_vec, comm_output_vec, done]() {
-      for (auto t : (*comm_input_vec)) {
-        delete t;
-      }
-      delete comm_input_vec;
-      for (auto t : (*comm_output_vec)) {
-        delete t;
-      }
-      delete comm_output_vec;
-      delete inputs_vec;
-      done();
-    };
-
-    comm->RunAsync(
-        "NcclAlltoallw", ctx, done_,
-        [inputs_vec, comm_input_vec, comm_output_vec, this, comm, ctx,
-         done_]() {
-          AllocatorAttributes host_alloc_attrs;
-          host_alloc_attrs.set_on_host(true);
-          host_alloc_attrs.set_gpu_compatible(true);
-
-          // Allocate scratch tensors.
-          Tensor* input_sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64, TensorShape({comm->size()}),
-                                 input_sizes, ctx->output_alloc_attr(0)),
-              done_);
-
-          Tensor host_input_sizes;
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64, TensorShape({comm->size()}),
-                                 &host_input_sizes, host_alloc_attrs),
-              done_);
-          for (int i = 0; i < comm->size(); ++i) {
-            host_input_sizes.flat<int64>()(i) = inputs_vec->at(i).NumElements();
-          }
-          ThenCopyToDevice(ctx, input_sizes, host_input_sizes);
-
-          Tensor* sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64,
-                                 TensorShape({comm->size() * comm->size()}),
-                                 sizes, ctx->output_alloc_attr(0)),
-              done_);
-
-          Tensor* host_sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64,
-                                 TensorShape({comm->size() * comm->size()}),
-                                 host_sizes, host_alloc_attrs),
-              done_);
-          auto sizes_ready = ThenRecordEvent(ctx);
-
-          comm->ThenWaitFor(sizes_ready);
-          comm->BlockHostUntilDone();
-
-          // Collect sizes of all inputs across devices.
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Allgather]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Allgather(*input_sizes, sizes),
-                               done_);
-          auto allgather_ready = comm->ThenRecordEvent();
-
-          ThenWaitFor(ctx, allgather_ready);
-          ThenCopyToHost(ctx, host_sizes, *sizes);
-          BlockHostUntilDone(ctx);
-          delete input_sizes;
-          delete sizes;
-
-          std::vector<int64> output_sizes_vec;
-          for (int64 i = 0; i < host_sizes->NumElements(); ++i) {
-            const int64 output_elements = host_sizes->flat<int64>()(i);
-            OP_REQUIRES_ASYNC(
-                ctx, output_elements % common_size_ == 0,
-                errors::InvalidArgument("common_shape size ", common_size_,
-                                        " is not compatible with input ", i,
-                                        ": ", output_elements),
-                done_);
-            const int64 output_size = output_elements / common_size_;
-            output_sizes_vec.push_back(output_size);
-          }
-          delete host_sizes;
-
-          // Allocate outputs for non-symmetric inputs.
-          for (int i = 0; i < comm->size(); ++i) {
-            if (i == comm->rank()) {
-              continue;
-            }
-            TensorShape output_shape;
-            PartialTensorShape(
-                {output_sizes_vec[comm->size() * i + comm->rank()]})
-                .Concatenate(common_shape_)
-                .AsTensorShape(&output_shape);
-            Tensor* output;
-            OP_REQUIRES_OK_ASYNC(
-                ctx, ctx->allocate_output(i, output_shape, &output), done_);
-          }
-          OpOutputList outputs;
-          OP_REQUIRES_OK_ASYNC(ctx, ctx->output_list("outputs", &outputs),
-                               done_);
-
-          // Exchange partial inputs.
-          std::vector<Tensor*> outputs_vec;
-          for (int i = 0; i < comm->size(); ++i) {
-            outputs_vec.push_back(outputs[i]);
-          }
-
-          // Cast and communicate.
-          auto call = functor::NcclAlltoallwCall<DTYPE, WIRE_DTYPE>();
-          OP_REQUIRES_OK_ASYNC(ctx,
-                               call(*inputs_vec, &outputs_vec, comm_input_vec,
-                                    comm_output_vec, ctx, comm, this),
-                               done_);
-        });
-  }
-
- private:
-  PartialTensorShape common_shape_;
-  int64 common_size_;
-};
-
-#define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
-  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoallw")                        \
-                              .Device(DEVICE_GPU)                        \
-                              .TypeConstraint<DTYPE>("dtype")            \
-                              .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
-                          NcclAlltoallwOp<DTYPE, WIRE_DTYPE>);
-TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
-#undef REGISTER_KERNEL
-#endif  // GOOGLE_CUDA
-
-#endif
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallw_n.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallw_n.cc
deleted file mode 100644
index e57c4eec..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_alltoallw_n.cc
+++ /dev/null
@@ -1,325 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#include <absl/strings/str_cat.h>
-
-#include <tensorflow/core/framework/common_shape_fns.h>
-#include <tensorflow/core/framework/op_kernel.h>
-#include <tensorflow/core/framework/shape_inference.h>
-
-#include <vector>
-
-#include "hybridbackend/tensorflow/common/cast.h"
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
-
-namespace tensorflow {
-namespace hybridbackend {
-
-#if HYBRIDBACKEND_NCCL
-
-REGISTER_OP("HbNcclAlltoallwN")
-    .Output("outputs: num_shards * dtype")
-    .Input("handle: resource")
-    .Input("inputs: num_shards * dtype")
-    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
-    .Attr("wire_dtype: {" TF_OP_NCCL_WIRE_DTYPE_LIST "}")
-    .Attr("num_shards: int >= 1 = 1")
-    .Attr("num_columns: int >= 1 = 1")
-    .Attr("common_shapes: list(shape)")
-    .SetIsStateful()
-    .SetShapeFn([](shape_inference::InferenceContext* c) {
-      int64 num_columns;
-      TF_RETURN_IF_ERROR(c->GetAttr("num_columns", &num_columns));
-      int64 comm_size = c->num_outputs() / num_columns;
-      std::vector<PartialTensorShape> common_shapes;
-      TF_RETURN_IF_ERROR(c->GetAttr("common_shapes", &common_shapes));
-      for (int64 i = 0; i < num_columns; ++i) {
-        shape_inference::ShapeHandle shape;
-        TF_RETURN_IF_ERROR(
-            c->MakeShapeFromPartialTensorShape(common_shapes[i], &shape));
-        TF_RETURN_IF_ERROR(c->Concatenate(
-            c->Vector(shape_inference::InferenceContext::kUnknownDim), shape,
-            &shape));
-        for (int64 dim = i * comm_size; dim < (i + 1) * comm_size; ++dim) {
-          c->set_output(dim, shape);
-        }
-      }
-      return Status::OK();
-    })
-    .Doc(R"doc(
-Grouped AllToAllw using a NCCL communicator.
-
-outputs: Rotated tensors for each device.
-handle: Handle of a NCCL communicator.
-inputs: Tensors to rotate for each device.
-)doc");
-
-#if GOOGLE_CUDA
-
-namespace functor {
-template <typename DTYPE, typename WIRE_DTYPE>
-struct NcclAlltoallwNCall {
-  Status operator()(const std::vector<Tensor>& inputs_vec,
-                    const int64 num_columns, std::vector<Tensor*>* outputs_vec,
-                    std::vector<Tensor*>* comm_input_vec,
-                    std::vector<Tensor*>* comm_output_vec, OpKernelContext* ctx,
-                    NcclComm* comm, NcclCommAsyncOp* comm_op) {
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [AlltoallwN]";
-    TF_RETURN_IF_ERROR(comm->AlltoallwN(inputs_vec, num_columns, outputs_vec));
-    return Status::OK();
-  }
-};
-
-template <>
-struct NcclAlltoallwNCall<float, Eigen::half> {
-  Status operator()(const std::vector<Tensor>& inputs_vec,
-                    const int64 num_columns, std::vector<Tensor*>* outputs_vec,
-                    std::vector<Tensor*>* comm_input_vec,
-                    std::vector<Tensor*>* comm_output_vec, OpKernelContext* ctx,
-                    NcclComm* comm, NcclCommAsyncOp* comm_op) {
-    for (int i = 0; i < inputs_vec.size(); ++i) {
-      TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_HALF, inputs_vec[i].shape(),
-                                            comm_input_vec->at(i)));
-      TF_RETURN_IF_ERROR(ctx->allocate_temp(
-          DT_HALF, outputs_vec->at(i)->shape(), comm_output_vec->at(i)));
-    }
-    comm->ThenWaitFor(comm_op->ThenRecordEvent(ctx));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastIn]";
-    for (int i = 0; i < inputs_vec.size(); i++) {
-      functor::Cast<float, Eigen::half>()(inputs_vec[i], comm_input_vec->at(i),
-                                          ctx, comm->stream());
-    }
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name()
-            << "] [AlltoallwN]";
-    TF_RETURN_IF_ERROR(
-        comm->AlltoallwN(*comm_input_vec, num_columns, comm_output_vec));
-    VLOG(1) << comm->DebugString() << " [" << comm_op->name() << "] [CastOut]";
-    for (int i = 0; i < comm_output_vec->size(); i++) {
-      functor::Cast<Eigen::half, float>()(
-          *(comm_output_vec->at(i)), outputs_vec->at(i), ctx, comm->stream());
-    }
-    return Status::OK();
-  }
-};
-
-}  // namespace functor
-
-template <typename DTYPE, typename WIRE_DTYPE>
-class NcclAlltoallwNOp : public NcclCommAsyncOp {
- public:
-  explicit NcclAlltoallwNOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_columns", &num_columns_));
-    std::vector<PartialTensorShape> common_shapes;
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("common_shapes", &common_shapes));
-
-    for (int64 c = 0; c < num_columns_; ++c) {
-      TensorShape output_shape;
-      PartialTensorShape({1})
-          .Concatenate(common_shapes[c])
-          .AsTensorShape(&output_shape);
-      int64 common_shape_size = 1;
-      for (int64 dim = 1; dim < output_shape.dims(); ++dim) {
-        common_shape_size *= output_shape.dim_size(dim);
-      }
-      output_shapes_.push_back(std::move(output_shape));
-      common_shape_sizes_.push_back(common_shape_size);
-    }
-  }
-
-  virtual void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                                    DoneCallback done) override {
-    // Get num_columns * comm_size inputs.
-    // e.g. [tensor0_for_gpu0, tensor0_for_gpu1,
-    //       tensor1_for_gpu0, tensor1_for_gpu1,
-    //       tensor2_for_gpu0, tensor2_for_gpu1]
-
-    AllocatorAttributes local_alloc_attrs;
-    local_alloc_attrs.set_on_host(true);
-    local_alloc_attrs.set_gpu_compatible(true);
-
-    OpInputList inputs;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("inputs", &inputs), done);
-
-    // Collect sizes of all inputs across devices.
-    const int64 inputs_size = inputs.size();
-    std::vector<Tensor>* inputs_vec = new std::vector<Tensor>();
-    for (int i = 0; i < inputs_size; ++i) {
-      inputs_vec->push_back(inputs[i]);
-    }
-    std::vector<Tensor*>* comm_input_vec = new std::vector<Tensor*>();
-    for (int i = 0; i < inputs_size; ++i) {
-      comm_input_vec->push_back(new Tensor());
-    }
-    std::vector<Tensor*>* comm_output_vec = new std::vector<Tensor*>();
-    for (int i = 0; i < inputs_size; ++i) {
-      comm_output_vec->push_back(new Tensor());
-    }
-    auto done_ = [inputs_vec, comm_input_vec, comm_output_vec, done]() {
-      for (auto t : (*comm_input_vec)) {
-        delete t;
-      }
-      delete comm_input_vec;
-      for (auto t : (*comm_output_vec)) {
-        delete t;
-      }
-      delete comm_output_vec;
-      delete inputs_vec;
-      done();
-    };
-
-    comm->RunAsync(
-        "NcclAlltoallwN", ctx, done_,
-        [inputs_vec, comm_input_vec, comm_output_vec, this, inputs_size, comm,
-         ctx, done_]() {
-          AllocatorAttributes local_alloc_attrs;
-          local_alloc_attrs.set_on_host(true);
-          local_alloc_attrs.set_gpu_compatible(true);
-
-          // Allocate scratch tensors.
-          Tensor* local_sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64, TensorShape({inputs_size}),
-                                 local_sizes, ctx->output_alloc_attr(0)),
-              done_);
-
-          Tensor host_local_sizes;
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64, TensorShape({inputs_size}),
-                                 &host_local_sizes, local_alloc_attrs),
-              done_);
-          for (int i = 0; i < inputs_size; ++i) {
-            host_local_sizes.flat<int64>()(i) = inputs_vec->at(i).NumElements();
-          }
-          ThenCopyToDevice(ctx, local_sizes, host_local_sizes);
-
-          Tensor* all_sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64,
-                                 TensorShape({inputs_size * comm->size()}),
-                                 all_sizes, ctx->output_alloc_attr(0)),
-              done_);
-
-          Tensor* host_all_sizes = new Tensor();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              ctx->allocate_temp(DT_INT64,
-                                 TensorShape({inputs_size * comm->size()}),
-                                 host_all_sizes, local_alloc_attrs),
-              done_);
-
-          comm->ThenWaitFor(ThenRecordEvent(ctx));
-          comm->BlockHostUntilDone();
-
-          // Collect sizes of all inputs across devices.
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Allgather]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Allgather(*local_sizes, all_sizes),
-                               done_);
-          ThenWaitFor(ctx, comm->ThenRecordEvent());
-          ThenCopyToHost(ctx, host_all_sizes, *all_sizes);
-          BlockHostUntilDone(ctx);
-          delete local_sizes;
-          delete all_sizes;
-
-          std::vector<int64> output_sizes_vec;
-          for (int64 i = 0; i < host_all_sizes->NumElements(); ++i) {
-            int64 to_idx = (i % inputs_size) % comm->size();
-            int64 c = (i % inputs_size) / comm->size();
-            int64 output_size = host_all_sizes->flat<int64>()(i);
-            int64 common_shape_size = common_shape_sizes_[c];
-            OP_REQUIRES_ASYNC(ctx, output_size % common_shape_size == 0,
-                              errors::InvalidArgument(absl::StrCat(
-                                  "[GPU", comm->rank(), "] Tensor ", c,
-                                  ": common_shape (size=", common_shape_size,
-                                  ") is not compatible with input ", to_idx,
-                                  " (size=", output_size, ")")),
-                              done_);
-            output_sizes_vec.push_back(output_size / common_shape_size);
-          }
-          delete host_all_sizes;
-
-          // Redirect symmetric inputs to outputs.
-          for (int64 c = 0; c < num_columns_; ++c) {
-            int64 fixed_idx = c * comm->size() + comm->rank();
-            ctx->set_output(fixed_idx, inputs_vec->at(fixed_idx));
-          }
-
-          // Allocate outputs for non-symmetric inputs.
-          for (int gpu_idx = 0; gpu_idx < comm->size(); ++gpu_idx) {
-            if (gpu_idx == comm->rank()) {
-              continue;
-            }
-            for (int64 c = 0; c < num_columns_; ++c) {
-              TensorShape output_shape(output_shapes_[c]);
-              output_shape.set_dim(
-                  0, output_sizes_vec[comm->size() * num_columns_ * gpu_idx +
-                                      comm->size() * c + comm->rank()]);
-              Tensor* output;
-              OP_REQUIRES_OK_ASYNC(
-                  ctx,
-                  ctx->allocate_output(comm->size() * c + gpu_idx, output_shape,
-                                       &output),
-                  done_);
-            }
-          }
-
-          OpOutputList outputs;
-          OP_REQUIRES_OK_ASYNC(ctx, ctx->output_list("outputs", &outputs),
-                               done_);
-
-          // Exchange partial inputs.
-          std::vector<Tensor*> outputs_vec;
-          for (int i = 0; i < outputs.size(); ++i) {
-            outputs_vec.push_back(outputs[i]);
-          }
-          // Cast and communicate.
-          auto call = functor::NcclAlltoallwNCall<DTYPE, WIRE_DTYPE>();
-          OP_REQUIRES_OK_ASYNC(
-              ctx,
-              call(*inputs_vec, num_columns_, &outputs_vec, comm_input_vec,
-                   comm_output_vec, ctx, comm, this),
-              done_);
-        });
-  }
-
- private:
-  int64 num_columns_;
-  std::vector<TensorShape> output_shapes_;
-  std::vector<int64> common_shape_sizes_;
-};
-
-#define REGISTER_KERNEL(DTYPE, WIRE_DTYPE)                               \
-  REGISTER_KERNEL_BUILDER(Name("HbNcclAlltoallwN")                       \
-                              .Device(DEVICE_GPU)                        \
-                              .TypeConstraint<DTYPE>("dtype")            \
-                              .TypeConstraint<WIRE_DTYPE>("wire_dtype"), \
-                          NcclAlltoallwNOp<DTYPE, WIRE_DTYPE>);
-TF_CALL_NCCL_CAST_TYPES(REGISTER_KERNEL);
-#undef REGISTER_KERNEL
-
-#endif  // GOOGLE_CUDA
-
-#endif
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_broadcast.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_broadcast.cc
index 41e73bad..d5d2e411 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_broadcast.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_broadcast.cc
@@ -21,7 +21,7 @@ limitations under the License.
 
 #include <vector>
 
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
 namespace hybridbackend {
@@ -29,11 +29,11 @@ namespace hybridbackend {
 #if HYBRIDBACKEND_NCCL
 
 REGISTER_OP("HbNcclBroadcast")
-    .Output("output: T")
+    .Output("output: dtype")
     .Input("handle: resource")
-    .Input("input: T")
+    .Input("input: dtype")
     .Attr("root_rank: int >= 0 = 0")
-    .Attr("T: {int8, uint8, int32, uint32, int64, uint64, half, float, double}")
+    .Attr("dtype: {" TF_OP_NCCL_DTYPE_LIST "}")
     .SetShapeFn([](shape_inference::InferenceContext* c) {
       c->set_output(0, c->input(1));
       return Status::OK();
@@ -49,18 +49,19 @@ root_rank: Rank of the broadcast root.
 )doc");
 
 #if GOOGLE_CUDA
-class NcclBroadcastOp : public NcclCommAsyncOp {
+class NcclBroadcastOp : public NcclCollectiveAsyncOp {
  public:
-  explicit NcclBroadcastOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {
+  explicit NcclBroadcastOp(OpKernelConstruction* ctx)
+      : NcclCollectiveAsyncOp(ctx) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("root_rank", &root_rank_));
     OP_REQUIRES(ctx, root_rank_ >= 0,
                 errors::InvalidArgument("root_rank should be >= 0"));
   }
 
-  void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                            DoneCallback done) override {
+  void CollectiveComputeAsync(NcclCollective* coll, OpKernelContext* ctx,
+                              DoneCallback done) override {
     OP_REQUIRES_ASYNC(
-        ctx, root_rank_ < comm->size(),
+        ctx, root_rank_ < coll->world_size(),
         errors::InvalidArgument("root_rank should be within communicator size"),
         done);
 
@@ -70,11 +71,12 @@ class NcclBroadcastOp : public NcclCommAsyncOp {
     OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, input->shape(), &output),
                          done);
 
-    comm->RunAsync(
-        "NcclBroadcast", ctx, done, [input, output, this, comm, ctx, done]() {
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Broadcast]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Broadcast(*input, root_rank_, output),
+    coll->stream()->LaunchUntilComputeDone(
+        ctx, [input, output, this, coll, ctx, done]() {
+          VLOG(1) << coll->DebugString() << " [" << name() << "] [Broadcast]";
+          OP_REQUIRES_OK_ASYNC(ctx, coll->Broadcast(*input, root_rank_, output),
                                done);
+          coll->stream()->BlockComputeUntilDone(ctx, done);
         });
   }
 
@@ -82,10 +84,11 @@ class NcclBroadcastOp : public NcclCommAsyncOp {
   int32 root_rank_;
 };
 
-#define REGISTER_KERNEL(TYPE)                                               \
-  REGISTER_KERNEL_BUILDER(                                                  \
-      Name("HbNcclBroadcast").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \
-      NcclBroadcastOp);
+#define REGISTER_KERNEL(TYPE)                                 \
+  REGISTER_KERNEL_BUILDER(Name("HbNcclBroadcast")             \
+                              .Device(DEVICE_GPU)             \
+                              .TypeConstraint<TYPE>("dtype"), \
+                          NcclBroadcastOp);
 TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 #endif  // GOOGLE_CUDA
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_collective.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_collective.cc
new file mode 100644
index 00000000..fde252d5
--- /dev/null
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_collective.cc
@@ -0,0 +1,484 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <mutex>
+#include <vector>
+
+#include <absl/strings/str_cat.h>
+
+#include <tensorflow/core/framework/common_shape_fns.h>
+#include <tensorflow/core/framework/op_kernel.h>
+#include <tensorflow/core/framework/shape_inference.h>
+
+#if GOOGLE_CUDA
+#include <tensorflow/stream_executor/cuda/cuda_activation.h>
+
+#include "hybridbackend/common/profiler.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
+#endif
+
+namespace tensorflow {
+namespace hybridbackend {
+
+#if HYBRIDBACKEND_NCCL
+#if GOOGLE_CUDA
+
+NcclCollective::NcclCollective(const string& shared_name, int world_size,
+                               int local_size, int rank)
+    : Collective(shared_name, world_size, local_size, rank),
+      stream_(new Stream) {}
+
+Status NcclCollective::Initialize(OpKernelContext* ctx) {
+  TF_RETURN_IF_ERROR(Collective::Initialize(ctx));
+
+  std::unique_lock<std::mutex> lock(mu_);
+  stream_->Initialize(
+      ctx, absl::StrCat("nccl_collective_", shared_name(), "_threads"),
+      3 /* num_threads */);
+  return Status::OK();
+}
+
+Status NcclCollective::Destroy() {
+  std::unique_lock<std::mutex> lock(mu_);
+  delete stream_;
+
+  if (!created_) {
+    return Status::OK();
+  }
+  created_ = false;
+  TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclCommDestroy(comm_)));
+  return Status::OK();
+}
+
+Status NcclCollective::Allreduce(const Tensor& input,
+                                 const CollectiveReduceOp reduce_op,
+                                 Tensor* output) {
+  const void* sendbuf = input.tensor_data().data();
+  void* recvbuf = const_cast<char*>(output->tensor_data().data());
+  const size_t count = input.NumElements();
+  if (TF_PREDICT_FALSE(count < 1)) {
+    return Status::OK();
+  }
+  ncclDataType_t nccl_dtype(ncclFloat);
+  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
+  ncclRedOp_t nccl_reduce_op(ncclSum);
+  TF_RETURN_IF_ERROR(ReduceOpToNcclReduceOp(reduce_op, &nccl_reduce_op));
+
+  std::unique_lock<std::mutex> lock(mu_);
+  return NcclErrorToStatus(ncclAllReduce(sendbuf, recvbuf, count, nccl_dtype,
+                                         nccl_reduce_op, comm_,
+                                         *(stream_->get())));
+}
+
+Status NcclCollective::AllreduceN(const std::vector<Tensor>& n_input,
+                                  const CollectiveReduceOp reduce_op,
+                                  std::vector<Tensor*>* n_output) {
+  std::unique_lock<std::mutex> lock(mu_);
+  ncclGroupStart();
+  for (size_t idx = 0; idx < n_input.size(); ++idx) {
+    const void* sendbuf = n_input[idx].tensor_data().data();
+    void* recvbuf = const_cast<char*>(n_output->at(idx)->tensor_data().data());
+    const size_t count = n_input[idx].NumElements();
+    if (TF_PREDICT_FALSE(count < 1)) {
+      continue;
+    }
+    ncclDataType_t nccl_dtype(ncclFloat);
+    TF_RETURN_IF_ERROR(EnumToNcclEnum(n_input[idx].dtype(), &nccl_dtype));
+    ncclRedOp_t nccl_reduce_op(ncclSum);
+    TF_RETURN_IF_ERROR(ReduceOpToNcclReduceOp(reduce_op, &nccl_reduce_op));
+
+    TF_RETURN_IF_ERROR(NcclErrorToStatus(
+        ncclAllReduce(sendbuf, recvbuf, count, nccl_dtype, nccl_reduce_op,
+                      comm_, *(stream_->get()))));
+  }
+  ncclGroupEnd();
+  return Status::OK();
+}
+
+Status NcclCollective::Alltoall(const Tensor& input, Tensor* output,
+                                CollectiveTopology topology) {
+  const size_t count = input.NumElements();
+  int active_size = Collective::compute_active_size(topology);
+  if (TF_PREDICT_FALSE(count == 0)) {
+    return Status::OK();
+  }
+  if (TF_PREDICT_FALSE(count % active_size != 0)) {
+    return errors::InvalidArgument("Number of elements in input (", count,
+                                   ") must can be divided into ", active_size,
+                                   " partitions");
+  }
+
+#if NCCL_VERSION_CODE >= 2700
+  const char* sendbuf = input.tensor_data().data();
+  char* recvbuf = const_cast<char*>(output->tensor_data().data());
+  const size_t partition_size = count / active_size;
+  ncclDataType_t nccl_dtype(ncclFloat);
+  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
+  const size_t dtype_size = DataTypeSize(input.dtype());
+
+  std::unique_lock<std::mutex> lock(mu_);
+  std::vector<int> active_ranks;
+  Collective::compute_active_ranks(topology, active_ranks);
+  ncclGroupStart();
+  for (int i = 0; i < active_ranks.size(); ++i) {
+    const size_t partition_bytes = i * partition_size * dtype_size;
+    TF_RETURN_IF_ERROR(NcclErrorToStatus(
+        ncclSend(sendbuf + partition_bytes, partition_size, nccl_dtype,
+                 active_ranks[i], comm_, *(stream_->get()))));
+    TF_RETURN_IF_ERROR(NcclErrorToStatus(
+        ncclRecv(recvbuf + partition_bytes, partition_size, nccl_dtype,
+                 active_ranks[i], comm_, *(stream_->get()))));
+  }
+  ncclGroupEnd();
+  return Status::OK();
+#else
+  return errors::Unimplemented("Alltoall not supported in NCCL < 2.7");
+#endif
+}
+
+Status NcclCollective::AlltoallN(const std::vector<Tensor>& n_input,
+                                 std::vector<Tensor*>* n_output,
+                                 CollectiveTopology topology) {
+  if (TF_PREDICT_FALSE(n_output->size() != n_input.size())) {
+    return errors::InvalidArgument(
+        "Inputs and outputs of AlltoallN must have same length");
+  }
+
+#if NCCL_VERSION_CODE >= 2700
+  std::unique_lock<std::mutex> lock(mu_);
+  int active_size = Collective::compute_active_size(topology);
+  std::vector<int> active_ranks;
+  compute_active_ranks(topology, active_ranks);
+  ncclGroupStart();
+  for (size_t idx = 0; idx < n_input.size(); ++idx) {
+    const size_t count = n_input[idx].NumElements();
+    if (TF_PREDICT_FALSE(count == 0)) {
+      continue;
+    }
+    if (TF_PREDICT_FALSE(count % active_size != 0)) {
+      return errors::InvalidArgument("Number of elements in input (", count,
+                                     ") must can be divided into ", active_size,
+                                     " partitions");
+    }
+    const char* sendbuf = n_input[idx].tensor_data().data();
+    char* recvbuf = const_cast<char*>(n_output->at(idx)->tensor_data().data());
+    const size_t partition_size = count / active_size;
+    ncclDataType_t nccl_dtype(ncclFloat);
+    TF_RETURN_IF_ERROR(EnumToNcclEnum(n_input[idx].dtype(), &nccl_dtype));
+    const size_t dtype_size = DataTypeSize(n_input[idx].dtype());
+
+    for (int i = 0; i < active_ranks.size(); ++i) {
+      const size_t partition_bytes = i * partition_size * dtype_size;
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclSend(sendbuf + partition_bytes, partition_size, nccl_dtype,
+                   active_ranks[i], comm_, *(stream_->get()))));
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclRecv(recvbuf + partition_bytes, partition_size, nccl_dtype,
+                   active_ranks[i], comm_, *(stream_->get()))));
+    }
+  }
+  ncclGroupEnd();
+  return Status::OK();
+#else
+  return errors::Unimplemented("AlltoallN not supported in NCCL < 2.7");
+#endif
+}
+
+Status NcclCollective::AlltoallN(const std::vector<Tensor*>& n_input,
+                                 std::vector<Tensor*>* n_output,
+                                 CollectiveTopology topology) {
+  if (TF_PREDICT_FALSE(n_output->size() != n_input.size())) {
+    return errors::InvalidArgument(
+        "Inputs and outputs of AlltoallN must have same length");
+  }
+
+#if NCCL_VERSION_CODE >= 2700
+  std::unique_lock<std::mutex> lock(mu_);
+  int active_size = Collective::compute_active_size(topology);
+  std::vector<int> active_ranks;
+  compute_active_ranks(topology, active_ranks);
+  ncclGroupStart();
+  for (size_t idx = 0; idx < n_input.size(); ++idx) {
+    const size_t count = n_input[idx]->NumElements();
+    if (TF_PREDICT_FALSE(count == 0)) {
+      continue;
+    }
+    if (TF_PREDICT_FALSE(count % active_size != 0)) {
+      return errors::InvalidArgument("Number of elements in input (", count,
+                                     ") must can be divided into ", active_size,
+                                     " partitions");
+    }
+    const char* sendbuf = n_input[idx]->tensor_data().data();
+    char* recvbuf = const_cast<char*>(n_output->at(idx)->tensor_data().data());
+    const size_t partition_size = count / active_size;
+    ncclDataType_t nccl_dtype(ncclFloat);
+    TF_RETURN_IF_ERROR(EnumToNcclEnum(n_input[idx]->dtype(), &nccl_dtype));
+    const size_t dtype_size = DataTypeSize(n_input[idx]->dtype());
+    ncclGroupStart();
+    for (int i = 0; i < active_ranks.size(); ++i) {
+      const size_t partition_bytes = i * partition_size * dtype_size;
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclSend(sendbuf + partition_bytes, partition_size, nccl_dtype,
+                   active_ranks[i], comm_, *(stream_->get()))));
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclRecv(recvbuf + partition_bytes, partition_size, nccl_dtype,
+                   active_ranks[i], comm_, *(stream_->get()))));
+    }
+    ncclGroupEnd();
+  }
+  ncclGroupEnd();
+  return Status::OK();
+#else
+  return errors::Unimplemented("AlltoallN not supported in NCCL < 2.7");
+#endif
+}
+
+Status NcclCollective::Alltoallv(const Tensor& input, const int32* send_sizes,
+                                 const int32* recv_sizes,
+                                 const int64 common_size, Tensor* output,
+                                 CollectiveTopology topology) {
+#if NCCL_VERSION_CODE >= 2700
+  const char* sendbuf = input.tensor_data().data();
+  char* recvbuf = const_cast<char*>(output->tensor_data().data());
+  ncclDataType_t nccl_dtype(ncclFloat);
+  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
+  const int32 dtype_size = static_cast<int32>(DataTypeSize(input.dtype()));
+
+  int32 sendoffset = 0;
+  int32 recvoffset = 0;
+  std::unique_lock<std::mutex> lock(mu_);
+  std::vector<int> active_ranks;
+  compute_active_ranks(topology, active_ranks);
+  ncclGroupStart();
+  for (int i = 0; i < active_ranks.size(); ++i) {
+    const int32 sendsize = send_sizes[i] * common_size;
+    const int32 recvsize = recv_sizes[i] * common_size;
+    if (TF_PREDICT_TRUE(sendsize > 0)) {
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclSend(sendbuf + sendoffset, sendsize, nccl_dtype, active_ranks[i],
+                   comm_, *(stream_->get()))));
+    }
+    if (TF_PREDICT_TRUE(recvsize > 0)) {
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclRecv(recvbuf + recvoffset, recvsize, nccl_dtype, active_ranks[i],
+                   comm_, *(stream_->get()))));
+    }
+    sendoffset += sendsize * dtype_size;
+    recvoffset += recvsize * dtype_size;
+  }
+  ncclGroupEnd();
+  return Status::OK();
+#else
+  return errors::Unimplemented("Alltoallv not supported in NCCL < 2.7");
+#endif
+}
+
+Status NcclCollective::AlltoallvN(const std::vector<Tensor>& n_input,
+                                  const std::vector<int32*>& n_send_sizes,
+                                  const std::vector<int32*>& n_recv_sizes,
+                                  const std::vector<int64>& n_common_size,
+                                  std::vector<Tensor*>* n_output,
+                                  CollectiveTopology topology) {
+  if (TF_PREDICT_FALSE(n_output->size() != n_input.size())) {
+    return errors::InvalidArgument(
+        "Inputs and outputs of AlltoallvN must have same length");
+  }
+
+#if NCCL_VERSION_CODE >= 2700
+  std::unique_lock<std::mutex> lock(mu_);
+  std::vector<int> active_ranks;
+  compute_active_ranks(topology, active_ranks);
+  ncclGroupStart();
+  for (size_t idx = 0; idx < n_input.size(); ++idx) {
+    const char* sendbuf = n_input[idx].tensor_data().data();
+    char* recvbuf = const_cast<char*>(n_output->at(idx)->tensor_data().data());
+    ncclDataType_t nccl_dtype(ncclFloat);
+    TF_RETURN_IF_ERROR(EnumToNcclEnum(n_input[idx].dtype(), &nccl_dtype));
+    const int32 dtype_size =
+        static_cast<int32>(DataTypeSize(n_input[idx].dtype()));
+
+    int32 sendoffset = 0;
+    int32 recvoffset = 0;
+    ncclGroupStart();
+    for (int i = 0; i < active_ranks.size(); ++i) {
+      const int32 sendsize = n_send_sizes[idx][i] * n_common_size[idx];
+      const int32 recvsize = n_recv_sizes[idx][i] * n_common_size[idx];
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclSend(sendbuf + sendoffset, sendsize, nccl_dtype, active_ranks[i],
+                   comm_, *(stream_->get()))));
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclRecv(recvbuf + recvoffset, recvsize, nccl_dtype, active_ranks[i],
+                   comm_, *(stream_->get()))));
+      sendoffset += sendsize * dtype_size;
+      recvoffset += recvsize * dtype_size;
+    }
+    ncclGroupEnd();
+  }
+  ncclGroupEnd();
+  return Status::OK();
+#else
+  return errors::Unimplemented("AlltoallvN not supported in NCCL < 2.7");
+#endif
+}
+
+Status NcclCollective::AlltoallvN(const std::vector<Tensor*>& n_input,
+                                  const std::vector<int32*>& n_send_sizes,
+                                  const std::vector<int32*>& n_recv_sizes,
+                                  const std::vector<int64>& n_common_size,
+                                  std::vector<Tensor*>* n_output,
+                                  CollectiveTopology topology) {
+  if (TF_PREDICT_FALSE(n_output->size() != n_input.size())) {
+    return errors::InvalidArgument(
+        "Inputs and outputs of AlltoallvN must have same length");
+  }
+
+#if NCCL_VERSION_CODE >= 2700
+  std::unique_lock<std::mutex> lock(mu_);
+  std::vector<int> active_ranks;
+  compute_active_ranks(topology, active_ranks);
+  ncclGroupStart();
+  for (size_t idx = 0; idx < n_input.size(); ++idx) {
+    const char* sendbuf = n_input[idx]->tensor_data().data();
+    char* recvbuf = const_cast<char*>(n_output->at(idx)->tensor_data().data());
+    ncclDataType_t nccl_dtype(ncclFloat);
+    TF_RETURN_IF_ERROR(EnumToNcclEnum(n_input[idx]->dtype(), &nccl_dtype));
+    const int32 dtype_size =
+        static_cast<int32>(DataTypeSize(n_input[idx]->dtype()));
+
+    int32 sendoffset = 0;
+    int32 recvoffset = 0;
+    ncclGroupStart();
+    for (int i = 0; i < active_ranks.size(); ++i) {
+      const int32 sendsize = n_send_sizes[idx][i] * n_common_size[idx];
+      const int32 recvsize = n_recv_sizes[idx][i] * n_common_size[idx];
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclSend(sendbuf + sendoffset, sendsize, nccl_dtype, active_ranks[i],
+                   comm_, *(stream_->get()))));
+      TF_RETURN_IF_ERROR(NcclErrorToStatus(
+          ncclRecv(recvbuf + recvoffset, recvsize, nccl_dtype, active_ranks[i],
+                   comm_, *(stream_->get()))));
+      sendoffset += sendsize * dtype_size;
+      recvoffset += recvsize * dtype_size;
+    }
+    ncclGroupEnd();
+  }
+  ncclGroupEnd();
+  return Status::OK();
+#else
+  return errors::Unimplemented("AlltoallvN not supported in NCCL < 2.7");
+#endif
+}
+
+Status NcclCollective::Broadcast(const Tensor& input, const int root_rank,
+                                 Tensor* output) {
+  const void* sendbuf = input.tensor_data().data();
+  void* recvbuf = const_cast<char*>(output->tensor_data().data());
+  const size_t count = input.NumElements();
+  ncclDataType_t nccl_dtype(ncclFloat);
+  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
+
+  std::unique_lock<std::mutex> lock(mu_);
+  return NcclErrorToStatus(ncclBroadcast(sendbuf, recvbuf, count, nccl_dtype,
+                                         root_rank, comm_, *(stream_->get())));
+}
+
+Status NcclCollective::Allgather(const Tensor& input, Tensor* output) {
+  const void* sendbuf = input.tensor_data().data();
+  void* recvbuf = const_cast<char*>(output->tensor_data().data());
+  const size_t count = input.NumElements();
+  ncclDataType_t nccl_dtype(ncclFloat);
+  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
+
+  std::unique_lock<std::mutex> lock(mu_);
+  return NcclErrorToStatus(ncclAllGather(sendbuf, recvbuf, count, nccl_dtype,
+                                         comm_, *(stream_->get())));
+}
+
+Status NcclCollective::Allgatherv(const Tensor& input, const Tensor& host_sizes,
+                                  Tensor* output) {
+  const void* sendbuf = input.tensor_data().data();
+  char* recvbuf = const_cast<char*>(output->tensor_data().data());
+  const int32* sendcounts = host_sizes.flat<int32>().data();
+  ncclDataType_t nccl_dtype(ncclFloat);
+  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
+  const size_t dtype_size = DataTypeSize(input.dtype());
+  size_t offset = 0;
+
+  std::unique_lock<std::mutex> lock(mu_);
+  ncclGroupStart();
+  for (int rank = 0; rank < world_size(); ++rank) {
+    TF_RETURN_IF_ERROR(NcclErrorToStatus(
+        ncclBroadcast(sendbuf, recvbuf + offset, sendcounts[rank], nccl_dtype,
+                      rank, comm_, *(stream_->get()))));
+    offset += sendcounts[rank] * dtype_size;
+  }
+  ncclGroupEnd();
+
+  return Status::OK();
+}
+
+Status NcclCollective::Create(const string& id) {
+  std::unique_lock<std::mutex> lock(mu_);
+  if (!TF_PREDICT_TRUE(id.size() == NCCL_UNIQUE_ID_BYTES)) {
+    return errors::InvalidArgument(
+        absl::StrCat("NCCL ID ", id.c_str(), " is invalid."));
+  }
+
+  ncclUniqueId nccl_id;
+  memcpy(nccl_id.internal, &id[0], NCCL_UNIQUE_ID_BYTES);
+  TF_RETURN_IF_ERROR(NcclErrorToStatus(
+      ncclCommInitRank(&comm_, world_size(), nccl_id, rank())));
+  created_ = true;
+  return Status::OK();
+}
+
+Status NcclCollective::CheckAsyncErrors() {
+  std::unique_lock<std::mutex> lock(mu_);
+  if (TF_PREDICT_FALSE(!created_)) {
+    return Status::OK();
+  }
+  ncclResult_t async_error;
+  TF_RETURN_IF_ERROR(
+      NcclErrorToStatus(ncclCommGetAsyncError(comm_, &async_error)));
+  if (TF_PREDICT_TRUE(async_error == ncclSuccess)) {
+    return Status::OK();
+  }
+
+  LOG(ERROR) << "NCCL communication aborted: "
+             << ncclGetErrorString(async_error);
+  TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclCommAbort(comm_)));
+  return NcclErrorToStatus(async_error);
+}
+
+NcclCollectiveAsyncOp::NcclCollectiveAsyncOp(OpKernelConstruction* ctx)
+    : AsyncOpKernel(ctx) {}
+
+void NcclCollectiveAsyncOp::ComputeAsync(OpKernelContext* ctx,
+                                         AsyncOpKernel::DoneCallback done) {
+  NcclCollective* coll = nullptr;
+  OP_REQUIRES_OK_ASYNC(ctx, LookupResource(ctx, HandleFromInput(ctx, 0), &coll),
+                       done);
+  CollectiveComputeAsync(coll, ctx, done);
+};
+
+#endif
+#endif
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_comm.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_comm.cc
deleted file mode 100644
index 66624400..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_comm.cc
+++ /dev/null
@@ -1,673 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#include <vector>
-
-#include <absl/strings/str_cat.h>
-
-#include <tensorflow/core/framework/common_shape_fns.h>
-#include <tensorflow/core/framework/op_kernel.h>
-#include <tensorflow/core/framework/shape_inference.h>
-
-#if GOOGLE_CUDA
-#include <tensorflow/stream_executor/cuda/cuda_activation.h>
-
-#include "hybridbackend/common/profiler.h"
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
-#endif
-
-namespace tensorflow {
-namespace hybridbackend {
-
-#if HYBRIDBACKEND_NCCL
-#if GOOGLE_CUDA
-
-NcclComm::NcclComm() {}
-
-NcclComm::~NcclComm() {
-  if (created_) {
-    Destroy();
-  }
-}
-
-Status NcclComm::Initialize(int size, int rank, const string& shared_name,
-                            OpKernelContext* ctx) {
-  if (!TF_PREDICT_TRUE(0 <= rank && rank < size)) {
-    return errors::InvalidArgument(
-        absl::StrCat("NCCL rank ", rank, " or size ", size, " is invalid."));
-  }
-
-  size_ = size;
-  rank_ = rank;
-
-  string thread_pool_name("nccl_comm_thread_");
-  for (size_t i = 0; i < shared_name.size(); ++i) {
-    const char ch = shared_name[i];
-    if ((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') ||
-        (ch >= '0' && ch <= '9') || ch == '_' || ch == '-') {
-      thread_pool_name += ch;
-    } else {
-      thread_pool_name += '_';
-    }
-  }
-  thread_pool_.reset(new thread::ThreadPool(
-      ctx->env(), ThreadOptions(), thread_pool_name, 1 /* num_threads */,
-      false /* low_latency_hint */));
-
-  ctx_stream_ = new se::Stream(ctx->op_device_context()->stream()->parent());
-  // NOTE(zycao): CU_STREAM_NON_BLOCKING was not supported yet.
-  // Using this argument for non-blocking by default stream.
-  ctx_stream_->Init();
-
-#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) <= 1010L
-  stream_ = reinterpret_cast<cudaStream_t*>(
-      ctx_stream_->implementation()->CudaStreamMemberHack());
-#else
-  stream_ = reinterpret_cast<cudaStream_t*>(
-      ctx_stream_->implementation()->GpuStreamMemberHack());
-#endif
-
-  debug_string_ = absl::StrCat("NcclComm(name=", shared_name, ", size=", size,
-                               ", rank=", rank, ")");
-
-  return Status::OK();
-}
-
-Status NcclComm::Create(const string& id) {
-  if (!TF_PREDICT_TRUE(id.size() == NCCL_UNIQUE_ID_BYTES)) {
-    return errors::InvalidArgument(
-        absl::StrCat("NCCL ID ", id.c_str(), " is invalid."));
-  }
-
-  ncclUniqueId nccl_id;
-  memcpy(nccl_id.internal, &id[0], NCCL_UNIQUE_ID_BYTES);
-  TF_RETURN_IF_ERROR(
-      NcclErrorToStatus(ncclCommInitRank(&comm_, size_, nccl_id, rank_)));
-  created_ = true;
-
-  return Status::OK();
-}
-
-Status NcclComm::Destroy() {
-  TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclCommDestroy(comm_)));
-  created_ = false;
-  return Status::OK();
-}
-
-Status NcclComm::Abort() {
-  TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclCommAbort(comm_)));
-  created_ = false;
-  return Status::OK();
-}
-
-Status NcclComm::UserRank(int* rank) {
-  return NcclErrorToStatus(ncclCommUserRank(comm_, rank));
-}
-
-void NcclComm::RunAsync(const string& message, OpKernelContext* ctx,
-                        AsyncOpKernel::DoneCallback done,
-                        std::function<void()> func) {
-  int device_id;
-  cudaGetDevice(&device_id);
-  se::Event* inputs_ready =
-      new se::Event(ctx->op_device_context()->stream()->parent());
-  inputs_ready->Init();
-  ctx->op_device_context()->stream()->ThenRecordEvent(inputs_ready);
-  thread_pool_->Schedule(
-      [device_id, inputs_ready, func, this, ctx, done, message]() {
-        cudaSetDevice(device_id);
-        se::cuda::ScopedActivateExecutorContext context(ctx_stream_->parent());
-        this->ThenWaitFor(inputs_ready);
-        auto* range = ::hybridbackend::ProfilerRange::forSynch(message);
-        func();
-        auto outputs_ready = this->ThenRecordEvent();
-        ctx->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute(
-            ctx->op_device_context()->stream(),
-            [outputs_ready, ctx, done, range]() {
-              ctx->op_device_context()->stream()->ThenWaitFor(outputs_ready);
-              delete outputs_ready;
-              delete range;
-              done();
-            });
-      });
-}
-
-se::Event* NcclComm::ThenRecordEvent() {
-  se::Event* ev = new se::Event(ctx_stream_->parent());
-  ev->Init();
-  ctx_stream_->ThenRecordEvent(ev);
-  return ev;
-}
-
-void NcclComm::ThenWaitFor(se::Event* ev) {
-  ctx_stream_->ThenWaitFor(ev);
-  delete ev;
-}
-
-void NcclComm::BlockHostUntilDone() { ctx_stream_->BlockHostUntilDone(); }
-
-Status NcclComm::Reduce(const Tensor& input, const ncclRedOp_t reduce_op,
-                        const int root_rank, Tensor* output) {
-  const void* sendbuf = input.tensor_data().data();
-  void* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const size_t count = input.NumElements();
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-
-  return NcclErrorToStatus(ncclReduce(sendbuf, recvbuf, count, nccl_dtype,
-                                      reduce_op, root_rank, comm_, *stream_));
-}
-
-Status NcclComm::ReduceScatter(const Tensor& input, const ncclRedOp_t reduce_op,
-                               Tensor* output) {
-  const void* sendbuf = input.tensor_data().data();
-  void* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const size_t count = output->NumElements();
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-
-  return NcclErrorToStatus(ncclReduceScatter(
-      sendbuf, recvbuf, count, nccl_dtype, reduce_op, comm_, *stream_));
-}
-
-Status NcclComm::Allreduce(const Tensor& input, const ncclRedOp_t reduce_op,
-                           Tensor* output) {
-  const void* sendbuf = input.tensor_data().data();
-  void* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const size_t count = input.NumElements();
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-
-  return NcclErrorToStatus(ncclAllReduce(sendbuf, recvbuf, count, nccl_dtype,
-                                         reduce_op, comm_, *stream_));
-}
-
-Status NcclComm::Broadcast(const Tensor& input, const int root_rank,
-                           Tensor* output) {
-  const void* sendbuf = input.tensor_data().data();
-  void* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const size_t count = input.NumElements();
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-
-  return NcclErrorToStatus(ncclBroadcast(sendbuf, recvbuf, count, nccl_dtype,
-                                         root_rank, comm_, *stream_));
-}
-
-Status NcclComm::Allgather(const Tensor& input, Tensor* output) {
-  const void* sendbuf = input.tensor_data().data();
-  void* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const size_t count = input.NumElements();
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-
-  return NcclErrorToStatus(
-      ncclAllGather(sendbuf, recvbuf, count, nccl_dtype, comm_, *stream_));
-}
-
-Status NcclComm::GroupAllgather(const std::vector<Tensor>& inputs,
-                                Tensor* output) {
-  ncclGroupStart();
-  size_t offset = 0;
-  for (size_t i = 0; i < inputs.size(); ++i) {
-    const void* sendbuf = inputs[i].tensor_data().data();
-    void* recvbuf = const_cast<char*>(output->tensor_data().data()) + offset;
-    const size_t count = inputs[i].NumElements();
-    ncclDataType_t nccl_dtype(ncclFloat);
-    TF_RETURN_IF_ERROR(EnumToNcclEnum(inputs[i].dtype(), &nccl_dtype));
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(
-        ncclAllGather(sendbuf, recvbuf, count, nccl_dtype, comm_, *stream_)));
-    offset += size_ * count * DataTypeSize(inputs[i].dtype());
-  }
-  ncclGroupEnd();
-  return Status::OK();
-}
-
-Status NcclComm::Allgatherv(const Tensor& input, const Tensor& host_sizes,
-                            Tensor* output) {
-  const void* sendbuf = input.tensor_data().data();
-  char* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const int32* sendcounts = host_sizes.flat<int32>().data();
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-  const size_t dtype_size = DataTypeSize(input.dtype());
-
-  size_t offset = 0;
-  ncclGroupStart();
-  for (int i = 0; i < size_; ++i) {
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(
-        ncclBroadcast(sendbuf, recvbuf + offset, sendcounts[i], nccl_dtype, i,
-                      comm_, *stream_)));
-    offset += sendcounts[i] * dtype_size;
-  }
-  ncclGroupEnd();
-  return Status::OK();
-}
-
-Status NcclComm::Alltoall(const Tensor& input, Tensor* output) {
-  const size_t count = input.NumElements();
-  if (TF_PREDICT_FALSE(count % size_ != 0)) {
-    return errors::InvalidArgument("Number of elements in input (", count,
-                                   ") must can be divided into ", size_,
-                                   " partitions");
-  }
-
-#if NCCL_VERSION_CODE >= 2700
-  const char* sendbuf = input.tensor_data().data();
-  char* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const size_t partition_size = count / size_;
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-  const size_t dtype_size = DataTypeSize(input.dtype());
-
-  ncclGroupStart();
-  for (int i = 0; i < size_; ++i) {
-    const size_t partition_bytes = i * partition_size * dtype_size;
-    TF_RETURN_IF_ERROR(
-        NcclErrorToStatus(ncclSend(sendbuf + partition_bytes, partition_size,
-                                   nccl_dtype, i, comm_, *stream_)));
-    TF_RETURN_IF_ERROR(
-        NcclErrorToStatus(ncclRecv(recvbuf + partition_bytes, partition_size,
-                                   nccl_dtype, i, comm_, *stream_)));
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("Alltoall not supported in NCCL < 2.7");
-#endif
-}
-
-Status NcclComm::Alltoallv(const Tensor& input, const Tensor& host_sizes,
-                           const int64 common_size, Tensor* output) {
-#if NCCL_VERSION_CODE >= 2700
-  const char* sendbuf = input.tensor_data().data();
-  char* recvbuf = const_cast<char*>(output->tensor_data().data());
-  const int32* sendcounts = host_sizes.flat<int32>().data();
-  ncclDataType_t nccl_dtype(ncclFloat);
-  TF_RETURN_IF_ERROR(EnumToNcclEnum(input.dtype(), &nccl_dtype));
-  const int32 dtype_size = static_cast<int32>(DataTypeSize(input.dtype()));
-
-  int32 sendoffset = 0;
-  int32 recvoffset = 0;
-  ncclGroupStart();
-  for (int i = 0; i < size_; ++i) {
-    const int32 sendsize = sendcounts[size_ * rank_ + i] * common_size;
-    const int32 recvsize = sendcounts[size_ * i + rank_] * common_size;
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclSend(
-        sendbuf + sendoffset, sendsize, nccl_dtype, i, comm_, *stream_)));
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclRecv(
-        recvbuf + recvoffset, recvsize, nccl_dtype, i, comm_, *stream_)));
-    sendoffset += sendsize * dtype_size;
-    recvoffset += recvsize * dtype_size;
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("Alltoallv not supported in NCCL < 2.7");
-#endif
-}
-
-Status NcclComm::AlltoallvN(const std::vector<Tensor>& inputs,
-                            const Tensor& host_sizes,
-                            const std::vector<int64>& common_sizes,
-                            std::vector<Tensor*>* outputs) {
-  if (TF_PREDICT_FALSE(outputs->size() != inputs.size() ||
-                       host_sizes.NumElements() !=
-                           inputs.size() * size_ * size_)) {
-    return errors::InvalidArgument(
-        "Size of inputs and size of outputs must be same and host_sizes must "
-        "have all sizes for all inputs across all devices");
-  }
-
-#if NCCL_VERSION_CODE >= 2700
-  ncclGroupStart();
-  for (size_t gid = 0; gid < inputs.size(); ++gid) {
-    const char* sendbuf = inputs[gid].tensor_data().data();
-    char* recvbuf = const_cast<char*>(outputs->at(gid)->tensor_data().data());
-    const int32* sendcounts =
-        host_sizes.flat<int32>().data() + gid * size_ * size_;
-    ncclDataType_t nccl_dtype(ncclFloat);
-    TF_RETURN_IF_ERROR(EnumToNcclEnum(inputs[gid].dtype(), &nccl_dtype));
-    const int32 dtype_size =
-        static_cast<int32>(DataTypeSize(inputs[gid].dtype()));
-
-    int32 sendoffset = 0;
-    int32 recvoffset = 0;
-    ncclGroupStart();
-    for (int i = 0; i < size_; ++i) {
-      const int32 sendsize = sendcounts[size_ * rank_ + i] * common_sizes[gid];
-      const int32 recvsize = sendcounts[size_ * i + rank_] * common_sizes[gid];
-      TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclSend(
-          sendbuf + sendoffset, sendsize, nccl_dtype, i, comm_, *stream_)));
-      TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclRecv(
-          recvbuf + recvoffset, recvsize, nccl_dtype, i, comm_, *stream_)));
-      sendoffset += sendsize * dtype_size;
-      recvoffset += recvsize * dtype_size;
-    }
-    ncclGroupEnd();
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("AlltoallvN not supported in NCCL < 2.7");
-#endif
-}
-
-Status NcclComm::AlltoallvN(const std::vector<Tensor*>& inputs,
-                            const Tensor& host_sizes,
-                            const std::vector<int64>& common_sizes,
-                            std::vector<Tensor*>* outputs) {
-  if (TF_PREDICT_FALSE(outputs->size() != inputs.size() ||
-                       host_sizes.NumElements() !=
-                           inputs.size() * size_ * size_)) {
-    return errors::InvalidArgument(
-        "Size of inputs and size of outputs must be same and host_sizes must "
-        "have all sizes for all inputs across all devices");
-  }
-
-#if NCCL_VERSION_CODE >= 2700
-  ncclGroupStart();
-  for (size_t gid = 0; gid < inputs.size(); ++gid) {
-    const char* sendbuf = inputs[gid]->tensor_data().data();
-    char* recvbuf = const_cast<char*>(outputs->at(gid)->tensor_data().data());
-    const int32* sendcounts =
-        host_sizes.flat<int32>().data() + gid * size_ * size_;
-    ncclDataType_t nccl_dtype(ncclFloat);
-    TF_RETURN_IF_ERROR(EnumToNcclEnum(inputs[gid]->dtype(), &nccl_dtype));
-    const int32 dtype_size =
-        static_cast<int32>(DataTypeSize(inputs[gid]->dtype()));
-
-    int32 sendoffset = 0;
-    int32 recvoffset = 0;
-    ncclGroupStart();
-    for (int i = 0; i < size_; ++i) {
-      const int32 sendsize = sendcounts[size_ * rank_ + i] * common_sizes[gid];
-      const int32 recvsize = sendcounts[size_ * i + rank_] * common_sizes[gid];
-      TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclSend(
-          sendbuf + sendoffset, sendsize, nccl_dtype, i, comm_, *stream_)));
-      TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclRecv(
-          recvbuf + recvoffset, recvsize, nccl_dtype, i, comm_, *stream_)));
-      sendoffset += sendsize * dtype_size;
-      recvoffset += recvsize * dtype_size;
-    }
-    ncclGroupEnd();
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("AlltoallvN not supported in NCCL < 2.7");
-#endif
-}
-
-Status NcclComm::Alltoallw(const std::vector<Tensor>& inputs,
-                           std::vector<Tensor*>* outputs) {
-  if (TF_PREDICT_FALSE(inputs.size() != size_ || outputs->size() != size_)) {
-    return errors::InvalidArgument(
-        "Size of inputs and outputs must be equal to the communicator size");
-  }
-  std::vector<const char*> sendbufs;
-  std::vector<char*> recvbufs;
-  std::vector<size_t> sendcounts;
-  std::vector<size_t> recvcounts;
-  std::vector<ncclDataType_t> senddtypes;
-  std::vector<ncclDataType_t> recvdtypes;
-  for (int i = 0; i < size_; ++i) {
-    sendbufs.push_back(inputs[i].tensor_data().data());
-    recvbufs.push_back(const_cast<char*>((*outputs)[i]->tensor_data().data()));
-    sendcounts.push_back(inputs[i].NumElements());
-    recvcounts.push_back((*outputs)[i]->NumElements());
-    ncclDataType_t senddtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum(inputs[i].dtype(), &senddtype));
-    senddtypes.push_back(senddtype);
-    ncclDataType_t recvdtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum((*outputs)[i]->dtype(), &recvdtype));
-    recvdtypes.push_back(recvdtype);
-  }
-
-#if NCCL_VERSION_CODE >= 2700
-  ncclGroupStart();
-  for (int i = 0; i < size_; ++i) {
-    if (rank_ == i) {
-      if (TF_PREDICT_TRUE(sendbufs[i] == recvbufs[i])) {
-        continue;
-      }
-    }
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclSend(
-        sendbufs[i], sendcounts[i], senddtypes[i], i, comm_, *stream_)));
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclRecv(
-        recvbufs[i], recvcounts[i], recvdtypes[i], i, comm_, *stream_)));
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("Alltoallw not supported in NCCL < 2.7");
-#endif
-}
-
-Status NcclComm::Alltoallw(const std::vector<Tensor*>& inputs,
-                           std::vector<Tensor*>* outputs) {
-  if (TF_PREDICT_FALSE(inputs.size() != size_ || outputs->size() != size_)) {
-    return errors::InvalidArgument(
-        "Size of inputs and outputs must be equal to the communicator size");
-  }
-  std::vector<const char*> sendbufs;
-  std::vector<char*> recvbufs;
-  std::vector<size_t> sendcounts;
-  std::vector<size_t> recvcounts;
-  std::vector<ncclDataType_t> senddtypes;
-  std::vector<ncclDataType_t> recvdtypes;
-  for (int i = 0; i < size_; ++i) {
-    sendbufs.push_back(inputs[i]->tensor_data().data());
-    recvbufs.push_back(const_cast<char*>((*outputs)[i]->tensor_data().data()));
-    sendcounts.push_back(inputs[i]->NumElements());
-    recvcounts.push_back((*outputs)[i]->NumElements());
-    ncclDataType_t senddtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum(inputs[i]->dtype(), &senddtype));
-    senddtypes.push_back(senddtype);
-    ncclDataType_t recvdtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum((*outputs)[i]->dtype(), &recvdtype));
-    recvdtypes.push_back(recvdtype);
-  }
-
-#if NCCL_VERSION_CODE >= 2700
-  ncclGroupStart();
-  for (int i = 0; i < size_; ++i) {
-    if (rank_ == i) {
-      if (TF_PREDICT_TRUE(sendbufs[i] == recvbufs[i])) {
-        continue;
-      }
-    }
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclSend(
-        sendbufs[i], sendcounts[i], senddtypes[i], i, comm_, *stream_)));
-    TF_RETURN_IF_ERROR(NcclErrorToStatus(ncclRecv(
-        recvbufs[i], recvcounts[i], recvdtypes[i], i, comm_, *stream_)));
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("Alltoallw not supported in NCCL < 2.7");
-#endif
-}
-
-Status NcclComm::AlltoallwN(const std::vector<Tensor>& inputs,
-                            const int64 num_columns,
-                            std::vector<Tensor*>* outputs) {
-  if (TF_PREDICT_FALSE(inputs.size() != num_columns * size_ ||
-                       outputs->size() != num_columns * size_)) {
-    return errors::InvalidArgument(
-        "Size of send_tensors and recv_tensors must be num_columns times of "
-        "size of the communicator");
-  }
-
-  std::vector<const char*> sendbufs;
-  std::vector<char*> recvbufs;
-  std::vector<size_t> sendcounts;
-  std::vector<size_t> recvcounts;
-  std::vector<ncclDataType_t> senddtypes;
-  std::vector<ncclDataType_t> recvdtypes;
-  for (int i = 0; i < inputs.size(); ++i) {
-    sendbufs.push_back(inputs[i].tensor_data().data());
-    recvbufs.push_back(const_cast<char*>((*outputs)[i]->tensor_data().data()));
-    sendcounts.push_back(inputs[i].NumElements());
-    recvcounts.push_back((*outputs)[i]->NumElements());
-    ncclDataType_t senddtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum(inputs[i].dtype(), &senddtype));
-    senddtypes.push_back(senddtype);
-    ncclDataType_t recvdtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum((*outputs)[i]->dtype(), &recvdtype));
-    recvdtypes.push_back(recvdtype);
-  }
-
-#if NCCL_VERSION_CODE >= 2700
-  ncclGroupStart();
-  for (int g = 0; g < num_columns; g++) {
-    ncclGroupStart();
-    for (int i = 0; i < size_; i++) {
-      int idx = g * size_ + i;
-      if (rank_ == i) {
-        if (TF_PREDICT_TRUE(sendbufs[idx] == recvbufs[idx])) {
-          continue;
-        }
-      }
-      TF_RETURN_IF_ERROR(
-          NcclErrorToStatus(ncclSend(sendbufs[idx], sendcounts[idx],
-                                     senddtypes[idx], i, comm_, *stream_)));
-      TF_RETURN_IF_ERROR(
-          NcclErrorToStatus(ncclRecv(recvbufs[idx], recvcounts[idx],
-                                     recvdtypes[idx], i, comm_, *stream_)));
-    }
-    ncclGroupEnd();
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("AlltoallwN not supported in NCCL < 2.7");
-#endif
-}
-
-Status NcclComm::AlltoallwN(const std::vector<Tensor*>& inputs,
-                            const int64 num_columns,
-                            std::vector<Tensor*>* outputs) {
-  if (TF_PREDICT_FALSE(inputs.size() != num_columns * size_ ||
-                       outputs->size() != num_columns * size_)) {
-    return errors::InvalidArgument(
-        "Size of send_tensors and recv_tensors must be num_columns times of "
-        "size of the communicator");
-  }
-
-  std::vector<const char*> sendbufs;
-  std::vector<char*> recvbufs;
-  std::vector<size_t> sendcounts;
-  std::vector<size_t> recvcounts;
-  std::vector<ncclDataType_t> senddtypes;
-  std::vector<ncclDataType_t> recvdtypes;
-  for (int i = 0; i < inputs.size(); ++i) {
-    sendbufs.push_back(inputs[i]->tensor_data().data());
-    recvbufs.push_back(const_cast<char*>((*outputs)[i]->tensor_data().data()));
-    sendcounts.push_back(inputs[i]->NumElements());
-    recvcounts.push_back((*outputs)[i]->NumElements());
-    ncclDataType_t senddtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum(inputs[i]->dtype(), &senddtype));
-    senddtypes.push_back(senddtype);
-    ncclDataType_t recvdtype;
-    TF_RETURN_IF_ERROR(EnumToNcclEnum((*outputs)[i]->dtype(), &recvdtype));
-    recvdtypes.push_back(recvdtype);
-  }
-
-#if NCCL_VERSION_CODE >= 2700
-  ncclGroupStart();
-  for (int g = 0; g < num_columns; g++) {
-    ncclGroupStart();
-    for (int i = 0; i < size_; i++) {
-      int idx = g * size_ + i;
-      if (rank_ == i) {
-        if (TF_PREDICT_TRUE(sendbufs[idx] == recvbufs[idx])) {
-          continue;
-        }
-      }
-      TF_RETURN_IF_ERROR(
-          NcclErrorToStatus(ncclSend(sendbufs[idx], sendcounts[idx],
-                                     senddtypes[idx], i, comm_, *stream_)));
-      TF_RETURN_IF_ERROR(
-          NcclErrorToStatus(ncclRecv(recvbufs[idx], recvcounts[idx],
-                                     recvdtypes[idx], i, comm_, *stream_)));
-    }
-    ncclGroupEnd();
-  }
-  ncclGroupEnd();
-  return Status::OK();
-#else
-  return errors::Unimplemented("AlltoallwN not supported in NCCL < 2.7");
-#endif
-}
-
-NcclCommAsyncOp::NcclCommAsyncOp(OpKernelConstruction* ctx)
-    : AsyncOpKernel(ctx) {}
-
-void NcclCommAsyncOp::ComputeAsync(OpKernelContext* ctx,
-                                   AsyncOpKernel::DoneCallback done) {
-  NcclComm* comm = nullptr;
-  OP_REQUIRES_OK_ASYNC(ctx, LookupResource(ctx, HandleFromInput(ctx, 0), &comm),
-                       done);
-  ComputeAsyncWithComm(comm, ctx, done);
-};
-
-se::Event* NcclCommAsyncOp::ThenRecordEvent(OpKernelContext* ctx) {
-  se::Event* ev = new se::Event(ctx->op_device_context()->stream()->parent());
-  ev->Init();
-  ctx->op_device_context()->stream()->ThenRecordEvent(ev);
-  return ev;
-}
-
-void NcclCommAsyncOp::ThenWaitFor(OpKernelContext* ctx, se::Event* ev) {
-  ctx->op_device_context()->stream()->ThenWaitFor(ev);
-  delete ev;
-}
-
-void NcclCommAsyncOp::ThenExecute(OpKernelContext* ctx,
-                                  std::function<void()> func) {
-  ctx->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute(
-      ctx->op_device_context()->stream(), func);
-}
-
-void NcclCommAsyncOp::ThenCopyToDevice(OpKernelContext* ctx, Tensor* dst,
-                                       const Tensor& src) {
-  se::DeviceMemoryBase dst_ptr(const_cast<char*>(dst->tensor_data().data()),
-                               dst->TotalBytes());
-  ctx->op_device_context()->stream()->ThenMemcpy(
-      &dst_ptr, src.tensor_data().data(), src.TotalBytes());
-}
-
-void NcclCommAsyncOp::ThenCopyToHost(OpKernelContext* ctx, Tensor* dst,
-                                     const Tensor& src) {
-  se::DeviceMemoryBase src_ptr(const_cast<char*>(src.tensor_data().data()),
-                               src.TotalBytes());
-  ctx->op_device_context()->stream()->ThenMemcpy(
-      const_cast<char*>(dst->tensor_data().data()), src_ptr, dst->TotalBytes());
-}
-
-void NcclCommAsyncOp::BlockHostUntilDone(OpKernelContext* ctx) {
-  ctx->op_device_context()->stream()->BlockHostUntilDone();
-}
-
-#endif
-#endif
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_create.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_create.cc
index 3ee4cdd3..48dc7c10 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_create.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_create.cc
@@ -21,16 +21,17 @@ limitations under the License.
 
 #include <vector>
 
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
 namespace hybridbackend {
 
 #if HYBRIDBACKEND_NCCL
 
-REGISTER_RESOURCE_HANDLE_OP(HbNcclComm);
+REGISTER_RESOURCE_HANDLE_OP(HbNcclCollective);
 
-REGISTER_OP("HbIsNcclCommInitialized")
+REGISTER_OP("HbIsNcclCollectiveInitialized")
     .Output("is_initialized: bool")
     .Input("handle: resource")
     .SetShapeFn(tensorflow::shape_inference::ScalarShape)
@@ -41,10 +42,11 @@ is_initialized: True if the NCCL communicator is initialized.
 handle: Handle of a NCCL communicator.
 )doc");
 
-REGISTER_OP("HbCreateNcclComm")
+REGISTER_OP("HbCreateNcclCollective")
     .Input("handle: resource")
     .Input("id: int64")
-    .Attr("size: int")
+    .Attr("world_size: int")
+    .Attr("local_size: int")
     .Attr("rank: int")
     .Attr("shared_name: string")
     .SetShapeFn(tensorflow::shape_inference::NoOutputs)
@@ -53,25 +55,28 @@ Creates a NCCL communicator and returns a handle to it.
 
 handle: Handle of a NCCL communicator.
 id: Unique ID of the NCCL communicator.
-size: Total number of ranks in the communicator.
+world_size: Total number of ranks in the communicator.
+local_size: Total number of ranks in a node.
 rank: Current rank in the communicator.
 shared_name: Shared name of all communicator instances.
 )doc");
 
 #if GOOGLE_CUDA
-REGISTER_KERNEL_BUILDER(Name("HbNcclCommHandleOp").Device(DEVICE_GPU),
-                        ResourceHandleOp<NcclComm>);
+REGISTER_KERNEL_BUILDER(Name("HbNcclCollectiveHandleOp").Device(DEVICE_GPU),
+                        ResourceHandleOp<NcclCollective>);
 
-REGISTER_KERNEL_BUILDER(Name("HbIsNcclCommInitialized")
+REGISTER_KERNEL_BUILDER(Name("HbIsNcclCollectiveInitialized")
                             .Device(DEVICE_GPU)
                             .HostMemory("is_initialized")
                             .HostMemory("handle"),
-                        IsResourceInitialized<NcclComm>);
+                        IsResourceInitialized<NcclCollective>);
 
-class CreateNcclCommOp : public AsyncOpKernel {
+class CreateNcclCollectiveOp : public AsyncOpKernel {
  public:
-  explicit CreateNcclCommOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) {
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("size", &size_));
+  explicit CreateNcclCollectiveOp(OpKernelConstruction* ctx)
+      : AsyncOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("world_size", &world_size_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("local_size", &local_size_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("rank", &rank_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("shared_name", &shared_name_));
   }
@@ -82,30 +87,48 @@ class CreateNcclCommOp : public AsyncOpKernel {
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("id", &id), done);
     string* nccl_id =
         new string(id->tensor_data().data(), NCCL_UNIQUE_ID_BYTES);
-    NcclComm* comm = new NcclComm();
-    OP_REQUIRES_OK_ASYNC(ctx, comm->Initialize(size_, rank_, shared_name_, ctx),
-                         done);
-
-    comm->RunAsync(
-        "NcclCommCreate", ctx, done, [this, comm, nccl_id, ctx, done]() {
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Create]";
-          OP_REQUIRES_OK_ASYNC(ctx, comm->Create(*nccl_id), done);
-          delete nccl_id;
-          Status s = CreateResource(ctx, HandleFromInput(ctx, 0), comm);
-          OP_REQUIRES_ASYNC(ctx, s.ok() || s.code() == error::ALREADY_EXISTS, s,
-                            done);
-        });
+    NcclCollective* coll =
+        new NcclCollective(shared_name_, world_size_, local_size_, rank_);
+    OP_REQUIRES_OK_ASYNC(ctx, coll->Initialize(ctx), done);
+
+    coll->stream()->LaunchUntilComputeDone(ctx, [this, coll, nccl_id, ctx,
+                                                 done]() {
+      VLOG(1) << coll->DebugString() << " [" << name() << "] [Create]";
+      OP_REQUIRES_OK_ASYNC(ctx, coll->Create(*nccl_id), done);
+      coll->stream()->BlockComputeUntilDone(ctx, [ctx, done, coll, nccl_id]() {
+        delete nccl_id;
+        Status s = CreateResource(ctx, HandleFromInput(ctx, 0), coll);
+        OP_REQUIRES_ASYNC(ctx, s.ok() || s.code() == error::ALREADY_EXISTS, s,
+                          done);
+        static const bool kNcclAsyncErrorHandling =
+            ::hybridbackend::EnvVarGetBool("NCCL_ASYNC_ERROR_HANDLING", true);
+        if (kNcclAsyncErrorHandling) {
+          coll->stream()->Launch(ctx, [ctx, done, coll]() {
+            static const bool kNcclAsyncErrorCheckingIntervalSecs =
+                ::hybridbackend::EnvVarGetInt(
+                    "NCCL_ASYNC_ERROR_CHECKING_INTERVAL_SECS", 10);
+            while (true) {
+              OP_REQUIRES_OK_ASYNC(ctx, coll->CheckAsyncErrors(), done);
+              Env::Default()->SleepForMicroseconds(
+                  kNcclAsyncErrorCheckingIntervalSecs * 1000000);
+            }
+          });
+        }
+        done();
+      });
+    });
   }
 
  private:
-  int size_;
+  int world_size_;
+  int local_size_;
   int rank_;
   string shared_name_;
 };
 
 REGISTER_KERNEL_BUILDER(
-    Name("HbCreateNcclComm").Device(DEVICE_GPU).HostMemory("id"),
-    CreateNcclCommOp);
+    Name("HbCreateNcclCollective").Device(DEVICE_GPU).HostMemory("id"),
+    CreateNcclCollectiveOp);
 #endif  // GOOGLE_CUDA
 
 #endif
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_get_id.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_get_id.cc
index 1f37d656..1b495792 100644
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_get_id.cc
+++ b/hybridbackend/tensorflow/distribute/nccl/nccl_get_id.cc
@@ -21,7 +21,7 @@ limitations under the License.
 
 #include <vector>
 
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
+#include "hybridbackend/tensorflow/distribute/nccl/collective.h"
 
 namespace tensorflow {
 namespace hybridbackend {
@@ -64,6 +64,7 @@ class GetNcclIdOp : public OpKernel {
 
 REGISTER_KERNEL_BUILDER(Name("HbGetNcclId").Device(DEVICE_GPU).HostMemory("id"),
                         GetNcclIdOp);
+REGISTER_KERNEL_BUILDER(Name("HbGetNcclId").Device(DEVICE_CPU), GetNcclIdOp);
 #endif
 
 #endif
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_reduce.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_reduce.cc
deleted file mode 100644
index 854cc846..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_reduce.cc
+++ /dev/null
@@ -1,106 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#include <tensorflow/core/framework/common_shape_fns.h>
-#include <tensorflow/core/framework/op_kernel.h>
-#include <tensorflow/core/framework/shape_inference.h>
-
-#include <vector>
-
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
-
-namespace tensorflow {
-namespace hybridbackend {
-
-#if HYBRIDBACKEND_NCCL
-
-REGISTER_OP("HbNcclReduce")
-    .Output("output: T")
-    .Input("handle: resource")
-    .Input("input: T")
-    .Attr("root_rank: int >= 0 = 0")
-    .Attr("reduce_op: int >= 0 = 0")
-    .Attr("T: {int8, uint8, int32, uint32, int64, uint64, half, float, double}")
-    .SetShapeFn([](shape_inference::InferenceContext* c) {
-      c->set_output(0, c->input(1));
-      return Status::OK();
-    })
-    .SetIsStateful()
-    .Doc(R"doc(
-Reduce using a NCCL communicator.
-
-output: A reduced tensor.
-handle: Handle of a NCCL communicator.
-input: A tensor to reduce.
-root_rank: Rank of the reduce root.
-reduce_op: Reduce ops: 0 for SUM, 1 for PROD, 2 for MAX, 3 for MIN.
-)doc");
-
-#if GOOGLE_CUDA
-class NcclReduceOp : public NcclCommAsyncOp {
- public:
-  explicit NcclReduceOp(OpKernelConstruction* ctx) : NcclCommAsyncOp(ctx) {
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("root_rank", &root_rank_));
-    OP_REQUIRES(ctx, root_rank_ >= 0,
-                errors::InvalidArgument("root_rank should be >= 0"));
-    int reduce_op;
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_op", &reduce_op));
-    OP_REQUIRES(ctx, reduce_op >= 0,
-                errors::InvalidArgument("reduce_op is invalid:", reduce_op));
-    OP_REQUIRES_OK(ctx, ReduceOpToNcclReduceOp(reduce_op, &reduce_op_));
-  }
-
-  void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                            DoneCallback done) override {
-    OP_REQUIRES_ASYNC(
-        ctx, root_rank_ < comm->size(),
-        errors::InvalidArgument("root_rank should be within communicator size"),
-        done);
-
-    const Tensor* input;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
-    Tensor* output;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, input->shape(), &output),
-                         done);
-
-    comm->RunAsync(
-        "NcclReduce", ctx, done, [input, output, this, comm, ctx, done]() {
-          VLOG(1) << comm->DebugString() << " [" << name() << "] [Reduce]";
-          OP_REQUIRES_OK_ASYNC(
-              ctx, comm->Reduce(*input, reduce_op_, root_rank_, output), done);
-        });
-  }
-
- private:
-  int32 root_rank_;
-  ncclRedOp_t reduce_op_;
-};
-
-#define REGISTER_KERNEL(TYPE)                                            \
-  REGISTER_KERNEL_BUILDER(                                               \
-      Name("HbNcclReduce").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \
-      NcclReduceOp);
-TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
-#undef REGISTER_KERNEL
-#endif  // GOOGLE_CUDA
-
-#endif
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/distribute/nccl/nccl_reduce_scatter.cc b/hybridbackend/tensorflow/distribute/nccl/nccl_reduce_scatter.cc
deleted file mode 100644
index 6a992b60..00000000
--- a/hybridbackend/tensorflow/distribute/nccl/nccl_reduce_scatter.cc
+++ /dev/null
@@ -1,125 +0,0 @@
-/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#if HYBRIDBACKEND_TENSORFLOW
-
-#include <tensorflow/core/framework/common_shape_fns.h>
-#include <tensorflow/core/framework/op_kernel.h>
-#include <tensorflow/core/framework/shape_inference.h>
-
-#include <vector>
-
-#include "hybridbackend/tensorflow/distribute/nccl/comm.h"
-
-namespace tensorflow {
-namespace hybridbackend {
-
-#if HYBRIDBACKEND_NCCL
-
-REGISTER_OP("HbNcclReduceScatter")
-    .Output("output: T")
-    .Input("handle: resource")
-    .Input("input: T")
-    .Attr("reduce_op: int >= 0 = 0")
-    .Attr("T: {int8, uint8, int32, uint32, int64, uint64, half, float, double}")
-    .SetShapeFn([](shape_inference::InferenceContext* c) {
-      shape_inference::ShapeHandle input = c->input(1);
-      if (!c->RankKnown(input)) {
-        c->set_output(0, c->UnknownShape());
-        return Status::OK();
-      }
-
-      int64 rank = c->Rank(input);
-      if (rank == 0) {
-        return errors::InvalidArgument(
-            "Scalar cannot be used in ReduceScatter communications.");
-      }
-      std::vector<shape_inference::DimensionHandle> dims(rank);
-      dims[0] = c->UnknownDim();
-      for (int32 i = 1; i < rank; ++i) {
-        dims[i] = c->Dim(input, i);
-      }
-      c->set_output(0, c->MakeShape(dims));
-      return Status::OK();
-    })
-    .SetIsStateful()
-    .Doc(R"doc(
-ReduceScatter using a NCCL communicator.
-
-output: A reduced and then scattered tensor.
-handle: Handle of a NCCL communicator.
-input: A tensor to reduce and scatter.
-reduce_op: Reduce ops: 0 for SUM, 1 for PROD, 2 for MAX, 3 for MIN.
-)doc");
-
-#if GOOGLE_CUDA
-class NcclReduceScatterOp : public NcclCommAsyncOp {
- public:
-  explicit NcclReduceScatterOp(OpKernelConstruction* ctx)
-      : NcclCommAsyncOp(ctx) {
-    int reduce_op;
-    OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_op", &reduce_op));
-    OP_REQUIRES(ctx, reduce_op >= 0,
-                errors::InvalidArgument("reduce_op is invalid:", reduce_op));
-    OP_REQUIRES_OK(ctx, ReduceOpToNcclReduceOp(reduce_op, &reduce_op_));
-  }
-
-  void ComputeAsyncWithComm(NcclComm* comm, OpKernelContext* ctx,
-                            DoneCallback done) override {
-    const Tensor* input;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->input("input", &input), done);
-    // Allocate output with a 1/N size buffer of input tensor, here N is the
-    // communicator size which indicates the device amounts.
-    TensorShape out_shape(input->shape());
-    OP_REQUIRES_ASYNC(
-        ctx, out_shape.dims() > 0 && out_shape.dim_size(0) % comm->size() == 0,
-        errors::InvalidArgument("Tensor cannot be scattered to ", comm->size(),
-                                " devices with shape ",
-                                out_shape.DebugString()),
-        done);
-    out_shape.set_dim(0, out_shape.dim_size(0) / comm->size());
-    Tensor* output;
-    OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, out_shape, &output),
-                         done);
-
-    comm->RunAsync("NcclReduceScatter", ctx, done,
-                   [input, output, this, comm, ctx, done]() {
-                     VLOG(1) << comm->DebugString() << " [" << name()
-                             << "] [ReduceScatter]";
-                     OP_REQUIRES_OK_ASYNC(
-                         ctx, comm->ReduceScatter(*input, reduce_op_, output),
-                         done);
-                   });
-  }
-
- private:
-  ncclRedOp_t reduce_op_;
-};
-
-#define REGISTER_KERNEL(TYPE)                             \
-  REGISTER_KERNEL_BUILDER(Name("HbNcclReduceScatter")     \
-                              .Device(DEVICE_GPU)         \
-                              .TypeConstraint<TYPE>("T"), \
-                          NcclReduceScatterOp);
-TF_CALL_NCCL_TYPES(REGISTER_KERNEL);
-#undef REGISTER_KERNEL
-#endif  // GOOGLE_CUDA
-
-#endif
-
-}  // namespace hybridbackend
-}  // namespace tensorflow
-
-#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/distribute/nccl/types.h b/hybridbackend/tensorflow/distribute/nccl/types.h
index 9ea48766..4e5ba3ea 100644
--- a/hybridbackend/tensorflow/distribute/nccl/types.h
+++ b/hybridbackend/tensorflow/distribute/nccl/types.h
@@ -28,6 +28,8 @@ limitations under the License.
 #include <tensorflow/core/framework/tensor.h>
 #include <tensorflow/core/lib/core/errors.h>
 
+#include "hybridbackend/tensorflow/distribute/collective.h"
+
 namespace tensorflow {
 namespace hybridbackend {
 
@@ -126,23 +128,23 @@ inline Status EnumToNcclEnum(const DataType& dtype,
   }
 }
 
-inline Status ReduceOpToNcclReduceOp(const int reduce_op,
+inline Status ReduceOpToNcclReduceOp(const CollectiveReduceOp reduce_op,
                                      ncclRedOp_t* nccl_reduce_op) {
   switch (reduce_op) {
-    case 0:
+    case kCollectiveSum:
       *nccl_reduce_op = ncclSum;
       return Status::OK();
-    case 1:
+    case kCollectiveProd:
       *nccl_reduce_op = ncclProd;
       return Status::OK();
-    case 2:
+    case kCollectiveMax:
       *nccl_reduce_op = ncclMax;
       return Status::OK();
-    case 3:
+    case kCollectiveMin:
       *nccl_reduce_op = ncclMin;
       return Status::OK();
 #if NCCL_VERSION_CODE >= 21000
-    case 4:
+    case kCollectiveAvg:
       *nccl_reduce_op = ncclAvg;
       return Status::OK();
 #endif
diff --git a/hybridbackend/tensorflow/distribute/ops.py b/hybridbackend/tensorflow/distribute/ops.py
new file mode 100644
index 00000000..eaf1128a
--- /dev/null
+++ b/hybridbackend/tensorflow/distribute/ops.py
@@ -0,0 +1,39 @@
+# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+r'''Collective constants.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+class CollectiveOps(object):  # pylint: disable=useless-object-inheritance
+  r'''Collective operations.
+  '''
+  SUM = 0
+  PROD = 1
+  MAX = 2
+  MIN = 3
+  AVG = 4
+
+
+class Topology(object):  # pylint: disable=useless-object-inheritance
+  r'''Communication topology.
+  '''
+  ALL = 0  # Communication across all GPUs
+  INTRA_NODE = 1  # Communication across all GPUs in current node
+  INTER_NODE = 2  # Communication across all GPUS with same rank in every nodes
diff --git a/hybridbackend/tensorflow/ops/floormod_shuffle/__init__.py b/hybridbackend/tensorflow/distribute/partition/__init__.py
similarity index 93%
rename from hybridbackend/tensorflow/ops/floormod_shuffle/__init__.py
rename to hybridbackend/tensorflow/distribute/partition/__init__.py
index dc05095b..a96f57f9 100644
--- a/hybridbackend/tensorflow/ops/floormod_shuffle/__init__.py
+++ b/hybridbackend/tensorflow/distribute/partition/__init__.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # =============================================================================
 
-r'''Floormod shuffle related classes and functions.
+r'''Partitioning related classes and functions.
 '''
 
 from __future__ import absolute_import
diff --git a/hybridbackend/tensorflow/ops/floormod_shuffle/functors.h b/hybridbackend/tensorflow/distribute/partition/modulo_functors.h
similarity index 84%
rename from hybridbackend/tensorflow/ops/floormod_shuffle/functors.h
rename to hybridbackend/tensorflow/distribute/partition/modulo_functors.h
index 8b863f8a..581e5ed6 100644
--- a/hybridbackend/tensorflow/ops/floormod_shuffle/functors.h
+++ b/hybridbackend/tensorflow/distribute/partition/modulo_functors.h
@@ -12,8 +12,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-#ifndef HYBRIDBACKEND_TENSORFLOW_EMBEDDING_FLOORMOD_SHUFFLE_FUNCTORS_H_
-#define HYBRIDBACKEND_TENSORFLOW_EMBEDDING_FLOORMOD_SHUFFLE_FUNCTORS_H_
+#ifndef HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_PARTITION_MODULO_FUNCTORS_H_
+#define HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_PARTITION_MODULO_FUNCTORS_H_
 
 #if HYBRIDBACKEND_TENSORFLOW
 
@@ -28,14 +28,14 @@ namespace hybridbackend {
 namespace functor {
 
 template <typename Device, typename T>
-struct FloormodShuffle {
+struct PartitionByModulo {
   void operator()(const int32 num_partitions, const Tensor& input,
                   Tensor* output, Tensor* sizes, Tensor* indices,
                   OpKernelContext* ctx);
 };
 
 template <typename Device, typename T>
-struct FloormodShuffleN {
+struct PartitionByModuloN {
   void operator()(const int32 num_partitions, const std::vector<Tensor>& inputs,
                   std::vector<Tensor*>& outputs,
                   std::vector<Tensor*>& outputs_sizes,
@@ -47,4 +47,4 @@ struct FloormodShuffleN {
 }  // namespace tensorflow
 
 #endif  // HYBRIDBACKEND_TENSORFLOW
-#endif  // HYBRIDBACKEND_TENSORFLOW_EMBEDDING_FLOORMOD_SHUFFLE_FUNCTORS_H_
+#endif  // HYBRIDBACKEND_TENSORFLOW_DISTRIBUTE_PARTITION_MODULO_FUNCTORS_H_
diff --git a/hybridbackend/tensorflow/ops/floormod_shuffle/ops.py b/hybridbackend/tensorflow/distribute/partition/ops.py
similarity index 67%
rename from hybridbackend/tensorflow/ops/floormod_shuffle/ops.py
rename to hybridbackend/tensorflow/distribute/partition/ops.py
index 49c3e8cc..bce8d99c 100644
--- a/hybridbackend/tensorflow/ops/floormod_shuffle/ops.py
+++ b/hybridbackend/tensorflow/distribute/partition/ops.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # =============================================================================
 
-r'''Arithmetic operators for dense and sparse tensors.
+r'''Partitioning message for collective ops.
 '''
 
 from __future__ import absolute_import
@@ -21,14 +21,16 @@
 from __future__ import print_function
 
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import data_flow_ops
 from tensorflow.python.ops import math_ops
 
-from hybridbackend.tensorflow.common import oplib as _ops
+from hybridbackend.tensorflow.framework.view import OperationLike
 
 
-def floormod_partition(ids, num_partitions):
+def partition(ids, num_partitions):
   r'''Partition IDs using floormod strategy.
 
    Args:
@@ -52,7 +54,7 @@ def floormod_partition(ids, num_partitions):
   return partitioned_ids, partitioned_indices
 
 
-def floormod_shuffle(ids, num_partitions, name=None):
+def _partition_by_modulo(ids, num_partitions, name=None):
   r'''Shuffle IDs using floormod strategy.
 
   Args:
@@ -66,31 +68,36 @@ def floormod_shuffle(ids, num_partitions, name=None):
     indices: Indices for gathering back.
   '''
   if name is None:
-    name = 'floormod_shuffle'
-
-  return _ops.hb_floormod_shuffle(
-    ids, num_partitions=num_partitions, name=name)
-
-
-def floormod_shuffle_n(ids_list, num_partitions, name=None):
-  r'''Shuffle multiple IDs using floormod strategy.
+    name = 'partition_by_modulo'
+
+  with ops.device(ids.device):
+    with ops.name_scope(name):
+      return (
+        OperationLike('PartitionByModulo')
+        .returns_tensors(
+          tensor_spec.TensorSpec(shape=ids.shape, dtype=ids.dtype),
+          tensor_spec.TensorSpec(shape=[num_partitions], dtype=dtypes.int32),
+          tensor_spec.TensorSpec(shape=ids.shape, dtype=dtypes.int32))
+        .finalize(
+          ids, num_partitions=num_partitions,
+          name=name))
+
+
+def partition_by_modulo(ids, num_partitions, name=None):
+  r'''Shuffle IDs using floormod strategy.
 
   Args:
-    ids_list: Input tensors to partition.
+    ids: Input tensor to partition.
     num_partitions: Number of partitions.
     name: Name of the operator.
 
   Return:
-    outputs: Tensors with shuffled IDs.
-    outputs_sizes: Size of each shard in outputs.
-    outputs_indices: Indices for gathering back to ids_list.
+    output: A tensor with shuffled IDs.
+    sizes: Size of each shard in output.
+    indices: Indices for gathering back.
   '''
-  if not isinstance(ids_list, (tuple, list)):
-    raise ValueError('ids_list must be a list')
-
   if name is None:
-    name = 'floormod_shuffle_n'
+    name = 'partition_by_modulo'
 
-  ids_list = list(ids_list)
-  return _ops.hb_floormod_shuffle_n(
-    ids_list, num_partitions=num_partitions, name=name)
+  return _partition_by_modulo(
+    ids, num_partitions=num_partitions, name=name)
diff --git a/hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_functors.cc b/hybridbackend/tensorflow/distribute/partition/partition_by_modulo_functors.cc
similarity index 79%
rename from hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_functors.cc
rename to hybridbackend/tensorflow/distribute/partition/partition_by_modulo_functors.cc
index 633ceea0..0e5fc44d 100644
--- a/hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_functors.cc
+++ b/hybridbackend/tensorflow/distribute/partition/partition_by_modulo_functors.cc
@@ -24,7 +24,7 @@ limitations under the License.
 #include <tensorflow/core/framework/tensor.h>
 #include <tensorflow/core/public/version.h>
 
-#include "hybridbackend/tensorflow/ops/floormod_shuffle/functors.h"
+#include "hybridbackend/tensorflow/distribute/partition/modulo_functors.h"
 
 namespace tensorflow {
 
@@ -35,7 +35,7 @@ namespace hybridbackend {
 namespace functor {
 
 template <typename T>
-struct FloormodShuffle<CPUDevice, T> {
+struct PartitionByModulo<CPUDevice, T> {
   void operator()(const int32 num_partitions, const Tensor& input,
                   Tensor* output, Tensor* sizes, Tensor* indices,
                   OpKernelContext* ctx) {
@@ -70,26 +70,26 @@ struct FloormodShuffle<CPUDevice, T> {
   }
 };
 
-template struct FloormodShuffle<CPUDevice, int32>;
-template struct FloormodShuffle<CPUDevice, int64>;
-template struct FloormodShuffle<CPUDevice, uint32>;
-template struct FloormodShuffle<CPUDevice, uint64>;
+template struct PartitionByModulo<CPUDevice, int32>;
+template struct PartitionByModulo<CPUDevice, int64>;
+template struct PartitionByModulo<CPUDevice, uint32>;
+template struct PartitionByModulo<CPUDevice, uint64>;
 
 template <typename T>
-struct FloormodShuffleN<CPUDevice, T> {
+struct PartitionByModuloN<CPUDevice, T> {
   void operator()(const int32 num_partitions, const std::vector<Tensor>& inputs,
                   std::vector<Tensor*>& outputs,
                   std::vector<Tensor*>& outputs_sizes,
                   std::vector<Tensor*>& outputs_indices, OpKernelContext* ctx) {
-    OP_REQUIRES_OK(
-        ctx, errors::Unimplemented("FloormodShuffleN on CPU not implemented."));
+    OP_REQUIRES_OK(ctx, errors::Unimplemented(
+                            "PartitionByModuloN on CPU not implemented."));
   }
 };
 
-template struct FloormodShuffleN<CPUDevice, int32>;
-template struct FloormodShuffleN<CPUDevice, int64>;
-template struct FloormodShuffleN<CPUDevice, uint32>;
-template struct FloormodShuffleN<CPUDevice, uint64>;
+template struct PartitionByModuloN<CPUDevice, int32>;
+template struct PartitionByModuloN<CPUDevice, int64>;
+template struct PartitionByModuloN<CPUDevice, uint32>;
+template struct PartitionByModuloN<CPUDevice, uint64>;
 
 }  // namespace functor
 }  // namespace hybridbackend
diff --git a/hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_functors.cu.cc b/hybridbackend/tensorflow/distribute/partition/partition_by_modulo_functors.cu.cc
similarity index 81%
rename from hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_functors.cu.cc
rename to hybridbackend/tensorflow/distribute/partition/partition_by_modulo_functors.cu.cc
index 97d151bd..a552293f 100644
--- a/hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_functors.cu.cc
+++ b/hybridbackend/tensorflow/distribute/partition/partition_by_modulo_functors.cu.cc
@@ -31,7 +31,7 @@ limitations under the License.
 #include "hybridbackend/common/profiler.h"
 
 #include "hybridbackend/tensorflow/common/device_functions.h"
-#include "hybridbackend/tensorflow/ops/floormod_shuffle/functors.h"
+#include "hybridbackend/tensorflow/distribute/partition/modulo_functors.h"
 
 namespace tensorflow {
 
@@ -42,10 +42,10 @@ namespace hybridbackend {
 namespace functor {
 
 template <typename T>
-__global__ void FloorModShuffleComputeSizes(const int32 num_partitions,
-                                            const int32 input_size,
-                                            const T* d_input, int32* d_sizes,
-                                            int32* d_offsets) {
+__global__ void PartitionByModuloComputeSizes(const int32 num_partitions,
+                                              const int32 input_size,
+                                              const T* d_input, int32* d_sizes,
+                                              int32* d_offsets) {
   for (int32 idx : CudaGridRangeX(input_size)) {
     const T v = d_input[idx];
     const T shard = (v % num_partitions + num_partitions) % num_partitions;
@@ -54,11 +54,12 @@ __global__ void FloorModShuffleComputeSizes(const int32 num_partitions,
 }
 
 template <typename T>
-__global__ void FloorModShufflePopulate(const int32 num_partitions,
-                                        const int32 input_size,
-                                        const T* d_input, const int32* d_sizes,
-                                        const int32* d_offsets, T* d_output,
-                                        int32* d_indices) {
+__global__ void PartitionByModuloPopulate(const int32 num_partitions,
+                                          const int32 input_size,
+                                          const T* d_input,
+                                          const int32* d_sizes,
+                                          const int32* d_offsets, T* d_output,
+                                          int32* d_indices) {
   extern __shared__ int32 outputs_offsets[];
   if (threadIdx.x == 0) {
     outputs_offsets[0] = 0;
@@ -78,7 +79,7 @@ __global__ void FloorModShufflePopulate(const int32 num_partitions,
 }
 
 template <typename T>
-struct FloormodShuffle<GPUDevice, T> {
+struct PartitionByModulo<GPUDevice, T> {
   void operator()(const int32 num_partitions, const Tensor& input,
                   Tensor* output, Tensor* sizes, Tensor* indices,
                   OpKernelContext* ctx) {
@@ -97,15 +98,15 @@ struct FloormodShuffle<GPUDevice, T> {
       return;
     }
 
-    auto* range = ::hybridbackend::ProfilerRange::forSynch("FloormodShuffle");
+    auto* range = ::hybridbackend::ProfilerRange::forSynch("PartitionByModulo");
     Tensor offsets_t;
     OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_INT32, TensorShape({input_size}),
                                            &offsets_t));
     int32* d_offsets = offsets_t.flat<int32>().data();
 
-    CudaLaunch(FloorModShuffleComputeSizes<T>, input_size, 0, d, nullptr,
+    CudaLaunch(PartitionByModuloComputeSizes<T>, input_size, 0, d, nullptr,
                num_partitions, input_size, d_input, d_sizes, d_offsets);
-    CudaLaunch(FloorModShufflePopulate<T>, input_size,
+    CudaLaunch(PartitionByModuloPopulate<T>, input_size,
                num_partitions * sizeof(int32), d, nullptr, num_partitions,
                input_size, d_input, d_sizes, d_offsets, d_output, d_indices);
     ctx->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute(
@@ -113,15 +114,15 @@ struct FloormodShuffle<GPUDevice, T> {
   }
 };
 
-template struct FloormodShuffle<GPUDevice, int32>;
-template struct FloormodShuffle<GPUDevice, int64>;
-template struct FloormodShuffle<GPUDevice, uint32>;
-template struct FloormodShuffle<GPUDevice, uint64>;
+template struct PartitionByModulo<GPUDevice, int32>;
+template struct PartitionByModulo<GPUDevice, int64>;
+template struct PartitionByModulo<GPUDevice, uint32>;
+template struct PartitionByModulo<GPUDevice, uint64>;
 
 template <typename T>
-__global__ void FloorModShuffleNZeroSizes(const int32 total_num_partitions,
-                                          const int32 num_partitions,
-                                          int32** hd_outputs_sizes) {
+__global__ void PartitionByModuloNZeroSizes(const int32 total_num_partitions,
+                                            const int32 num_partitions,
+                                            int32** hd_outputs_sizes) {
   for (int32 idx : CudaGridRangeX(total_num_partitions)) {
     const int32 s = idx / num_partitions;
     const int32 sidx = idx % num_partitions;
@@ -130,13 +131,11 @@ __global__ void FloorModShuffleNZeroSizes(const int32 total_num_partitions,
 }
 
 template <typename T>
-__global__ void FloorModShuffleNComputeSizes(const int32 total_max_inputs_size,
-                                             const int32 max_inputs_size,
-                                             const int32 num_partitions,
-                                             const int32* d_inputs_size,
-                                             const size_t* dd_inputs,
-                                             size_t* dd_outputs_sizes,
-                                             int32* d_output_segment_offsets) {
+__global__ void PartitionByModuloNComputeSizes(
+    const int32 total_max_inputs_size, const int32 max_inputs_size,
+    const int32 num_partitions, const int32* d_inputs_size,
+    const size_t* dd_inputs, size_t* dd_outputs_sizes,
+    int32* d_output_segment_offsets) {
   for (int32 idx : CudaGridRangeX(total_max_inputs_size)) {
     const int32 s = idx / max_inputs_size;
     const int32 sidx = idx % max_inputs_size;
@@ -150,10 +149,10 @@ __global__ void FloorModShuffleNComputeSizes(const int32 total_max_inputs_size,
   }
 }
 
-__global__ void FloorModShuffleNComputeOffsets(const int32 num_inputs,
-                                               const int32 num_partitions,
-                                               size_t* dd_outputs_sizes,
-                                               int32* d_output_offsets) {
+__global__ void PartitionByModuloNComputeOffsets(const int32 num_inputs,
+                                                 const int32 num_partitions,
+                                                 size_t* dd_outputs_sizes,
+                                                 int32* d_output_offsets) {
   for (int32 idx : CudaGridRangeX(num_inputs)) {
     int32* d_sizes = reinterpret_cast<int32*>(ldg(dd_outputs_sizes + idx));
     d_output_offsets[num_partitions * idx] = 0;
@@ -165,7 +164,7 @@ __global__ void FloorModShuffleNComputeOffsets(const int32 num_inputs,
 }
 
 template <typename T>
-__global__ void FloorModShuffleNPopulate(
+__global__ void PartitionByModuloNPopulate(
     const int32 total_max_inputs_size, const int32 max_inputs_size,
     const int32 num_partitions, const int32* d_inputs_size,
     const size_t* dd_inputs, size_t* dd_outputs, size_t* dd_outputs_indices,
@@ -188,12 +187,13 @@ __global__ void FloorModShuffleNPopulate(
 }
 
 template <typename T>
-struct FloormodShuffleN<GPUDevice, T> {
+struct PartitionByModuloN<GPUDevice, T> {
   void operator()(const int32 num_partitions, const std::vector<Tensor>& inputs,
                   std::vector<Tensor*>& outputs,
                   std::vector<Tensor*>& outputs_sizes,
                   std::vector<Tensor*>& outputs_indices, OpKernelContext* ctx) {
-    auto* range = ::hybridbackend::ProfilerRange::forSynch("FloormodShuffleN");
+    auto* range =
+        ::hybridbackend::ProfilerRange::forSynch("PartitionByModuloN");
     const int32 num_inputs = inputs.size();
     std::vector<int32> inputs_size;
     for (int i = 0; i < num_inputs; ++i) {
@@ -230,7 +230,7 @@ struct FloormodShuffleN<GPUDevice, T> {
     }
 
     const int32 total_num_partitions = num_inputs * num_partitions;
-    CudaLaunch(FloorModShuffleNZeroSizes<T>, total_num_partitions, 0, d,
+    CudaLaunch(PartitionByModuloNZeroSizes<T>, total_num_partitions, 0, d,
                nullptr, total_num_partitions, num_partitions,
                hd_all_outputs_sizes);
 
@@ -315,14 +315,14 @@ struct FloormodShuffleN<GPUDevice, T> {
     int32* d_output_offsets =
         d_offsets_buffer_t.flat<int32>().data() + total_max_inputs_size;
 
-    CudaLaunch(FloorModShuffleNComputeSizes<T>, total_max_inputs_size, 0, d,
+    CudaLaunch(PartitionByModuloNComputeSizes<T>, total_max_inputs_size, 0, d,
                nullptr, total_max_inputs_size, max_inputs_size, num_partitions,
                d_inputs_size, dd_inputs, dd_outputs_sizes,
                d_output_segment_offsets);
-    CudaLaunch(FloorModShuffleNComputeOffsets, num_nonzero_inputs, 0, d,
+    CudaLaunch(PartitionByModuloNComputeOffsets, num_nonzero_inputs, 0, d,
                nullptr, num_nonzero_inputs, num_partitions, dd_outputs_sizes,
                d_output_offsets);
-    CudaLaunch(FloorModShuffleNPopulate<T>, total_max_inputs_size, 0, d,
+    CudaLaunch(PartitionByModuloNPopulate<T>, total_max_inputs_size, 0, d,
                nullptr, total_max_inputs_size, max_inputs_size, num_partitions,
                d_inputs_size, dd_inputs, dd_outputs, dd_outputs_indices,
                d_output_segment_offsets, d_output_offsets);
@@ -349,10 +349,10 @@ struct FloormodShuffleN<GPUDevice, T> {
   }
 };
 
-template struct FloormodShuffleN<GPUDevice, int32>;
-template struct FloormodShuffleN<GPUDevice, int64>;
-template struct FloormodShuffleN<GPUDevice, uint32>;
-template struct FloormodShuffleN<GPUDevice, uint64>;
+template struct PartitionByModuloN<GPUDevice, int32>;
+template struct PartitionByModuloN<GPUDevice, int64>;
+template struct PartitionByModuloN<GPUDevice, uint32>;
+template struct PartitionByModuloN<GPUDevice, uint64>;
 
 }  // namespace functor
 }  // namespace hybridbackend
diff --git a/hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_ops.cc b/hybridbackend/tensorflow/distribute/partition/partition_by_modulo_ops.cc
similarity index 75%
rename from hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_ops.cc
rename to hybridbackend/tensorflow/distribute/partition/partition_by_modulo_ops.cc
index 985fb115..a4bd9abd 100644
--- a/hybridbackend/tensorflow/ops/floormod_shuffle/floormod_shuffle_ops.cc
+++ b/hybridbackend/tensorflow/distribute/partition/partition_by_modulo_ops.cc
@@ -28,7 +28,7 @@ limitations under the License.
 #include <tensorflow/core/framework/tensor.pb.h>
 
 #include "hybridbackend/common/env.h"
-#include "hybridbackend/tensorflow/ops/floormod_shuffle/functors.h"
+#include "hybridbackend/tensorflow/distribute/partition/modulo_functors.h"
 
 #if GOOGLE_CUDA
 #include "hybridbackend/tensorflow/common/host_functions.h"
@@ -37,11 +37,13 @@ limitations under the License.
 namespace tensorflow {
 
 using CPUDevice = Eigen::ThreadPoolDevice;
+#if GOOGLE_CUDA
 using GPUDevice = Eigen::GpuDevice;
+#endif  // GOOGLE_CUDA
 
 namespace hybridbackend {
 
-REGISTER_OP("HbFloormodShuffle")
+REGISTER_OP("HbPartitionByModulo")
     .Output("output: T")
     .Output("sizes: int32")
     .Output("indices: int32")
@@ -67,9 +69,9 @@ num_partitions: Number of partitions.
 )doc");
 
 template <typename Device, typename T>
-class FloormodShuffleOp : public OpKernel {
+class PartitionByModuloOp : public OpKernel {
  public:
-  FloormodShuffleOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
+  PartitionByModuloOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("num_partitions", &num_partitions_));
   }
 
@@ -78,7 +80,7 @@ class FloormodShuffleOp : public OpKernel {
 
     OP_REQUIRES(
         ctx, TensorShapeUtils::IsVector(input.shape()),
-        errors::InvalidArgument("floormod_shuffle expects a 1D vector."));
+        errors::InvalidArgument("partition_by_modulo expects a 1D vector."));
 
     const int32 input_size = input.NumElements();
     Tensor* output = nullptr;
@@ -90,8 +92,8 @@ class FloormodShuffleOp : public OpKernel {
     Tensor* indices = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(2, {input_size}, &indices));
 
-    functor::FloormodShuffle<Device, T>()(num_partitions_, input, output, sizes,
-                                          indices, ctx);
+    functor::PartitionByModulo<Device, T>()(num_partitions_, input, output,
+                                            sizes, indices, ctx);
   }
 
  private:
@@ -101,14 +103,25 @@ class FloormodShuffleOp : public OpKernel {
 #define TF_CALL_SHUFFLE_PARTITION_TYPES(m) \
   TF_CALL_int32(m) TF_CALL_uint32(m) TF_CALL_int64(m) TF_CALL_uint64(m)
 
-#define REGISTER_SHUFFLE_PARTITION_KERNEL(TYPE)                               \
-  REGISTER_KERNEL_BUILDER(                                                    \
-      Name("HbFloormodShuffle").Device(DEVICE_CPU).TypeConstraint<TYPE>("T"), \
-      FloormodShuffleOp<CPUDevice, TYPE>)
+#define REGISTER_SHUFFLE_PARTITION_KERNEL(TYPE)           \
+  REGISTER_KERNEL_BUILDER(Name("HbPartitionByModulo")     \
+                              .Device(DEVICE_CPU)         \
+                              .TypeConstraint<TYPE>("T"), \
+                          PartitionByModuloOp<CPUDevice, TYPE>)
 TF_CALL_SHUFFLE_PARTITION_TYPES(REGISTER_SHUFFLE_PARTITION_KERNEL);
 #undef REGISTER_SHUFFLE_PARTITION_KERNEL
 
-REGISTER_OP("HbFloormodShuffleN")
+#if GOOGLE_CUDA
+#define REGISTER_SHUFFLE_PARTITION_KERNEL(TYPE)           \
+  REGISTER_KERNEL_BUILDER(Name("HbPartitionByModulo")     \
+                              .Device(DEVICE_GPU)         \
+                              .TypeConstraint<TYPE>("T"), \
+                          PartitionByModuloOp<GPUDevice, TYPE>)
+TF_CALL_SHUFFLE_PARTITION_TYPES(REGISTER_SHUFFLE_PARTITION_KERNEL);
+#undef REGISTER_SHUFFLE_PARTITION_KERNEL
+#endif  // GOOGLE_CUDA
+
+REGISTER_OP("HbPartitionByModuloN")
     .Output("outputs: N * T")
     .Output("outputs_sizes: N * int32")
     .Output("outputs_indices: N * int32")
@@ -139,9 +152,9 @@ num_partitions: Number of partitions.
 )doc");
 
 template <typename Device, typename T>
-class FloormodShuffleNOp : public OpKernel {
+class PartitionByModuloNOp : public OpKernel {
  public:
-  FloormodShuffleNOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
+  PartitionByModuloNOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &num_inputs_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("num_partitions", &num_partitions_));
   }
@@ -158,7 +171,7 @@ class FloormodShuffleNOp : public OpKernel {
       const Tensor& input = inputs[i];
       OP_REQUIRES(ctx, TensorShapeUtils::IsVector(input.shape()),
                   errors::InvalidArgument(
-                      "floormod_shuffle_n expects 1D vector for input ", i));
+                      "partition_by_modulo_n expects 1D vector for input ", i));
       const int32 input_size = input.NumElements();
       inputs_vec.push_back(input);
       Tensor* output = nullptr;
@@ -174,7 +187,7 @@ class FloormodShuffleNOp : public OpKernel {
       indices_vec.push_back(indices);
     }
 
-    functor::FloormodShuffleN<Device, T>()(
+    functor::PartitionByModuloN<Device, T>()(
         num_partitions_, inputs_vec, outputs_vec, sizes_vec, indices_vec, ctx);
   }
 
@@ -184,10 +197,11 @@ class FloormodShuffleNOp : public OpKernel {
 };
 
 #if GOOGLE_CUDA
-#define REGISTER_SHUFFLE_PARTITION_KERNEL(TYPE)                                \
-  REGISTER_KERNEL_BUILDER(                                                     \
-      Name("HbFloormodShuffleN").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \
-      FloormodShuffleNOp<GPUDevice, TYPE>)
+#define REGISTER_SHUFFLE_PARTITION_KERNEL(TYPE)           \
+  REGISTER_KERNEL_BUILDER(Name("HbPartitionByModuloN")    \
+                              .Device(DEVICE_GPU)         \
+                              .TypeConstraint<TYPE>("T"), \
+                          PartitionByModuloNOp<GPUDevice, TYPE>)
 TF_CALL_SHUFFLE_PARTITION_TYPES(REGISTER_SHUFFLE_PARTITION_KERNEL);
 #undef REGISTER_SHUFFLE_PARTITION_KERNEL
 #endif  // GOOGLE_CUDA
diff --git a/hybridbackend/tensorflow/distribute/pubsub.py b/hybridbackend/tensorflow/distribute/pubsub.py
deleted file mode 100644
index 72da6e2d..00000000
--- a/hybridbackend/tensorflow/distribute/pubsub.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Classes for communication pub/sub.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import threading
-
-from tensorflow.python.framework import device as pydev
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import collective_ops
-
-from hybridbackend.tensorflow.framework.context import Context
-
-
-class PubSub(object):  # pylint: disable=useless-object-inheritance
-  r'''Publish/subscribe messages.
-  '''
-  _lock = threading.Lock()
-  _thread_local = threading.local()
-
-  _prev_channel = 1
-  _devices_to_channel = {}
-
-  _prev_call = 1000000
-  _name_to_call = {}
-
-  def __init__(self, devices, rank=None, root_rank=None):
-    r'''Constructs a pub/sub instance.
-
-    Args:
-      devices: Involved devices.
-      rank: Rank of this peer.
-      root_rank: Rank of the root peer.
-    '''
-    self._bcast_device = Context.get().options.comm_pubsub_device
-    self._devices = devices
-    parsed_devices = [pydev.DeviceSpec.from_string(d) for d in devices]
-    sorted_devices = sorted(
-      [f'{d.device_type}:{d.device_index}' for d in parsed_devices])
-    devices_str = ','.join(sorted_devices)
-    with PubSub._lock:
-      if devices_str not in PubSub._devices_to_channel:
-        channel = PubSub._prev_channel
-        PubSub._prev_channel += 1
-        PubSub._devices_to_channel[devices_str] = channel
-    self._channel = PubSub._devices_to_channel[devices_str]
-
-    if rank is None:
-      device = Context.current_device()
-      if device not in self._devices:
-        raise ValueError(
-          f'Current device {device} not in devices {self._devices}')
-      rank = self._devices.index(device)
-    self._rank = rank
-    if root_rank is None:
-      root_rank = 0
-    self._root_rank = root_rank
-
-  @property
-  def channel(self):
-    r'''Channel for pub/sub.
-    '''
-    return self._channel
-
-  def __call__(self, fn, shape, dtype, name=None):
-    r'''Publish/subscribe message across devices.
-
-    Args:
-      fn: Function to generate tensor.
-      shape: Shape of the generate tensor.
-      dtype: Data type of the generated tensor.
-      name: Name of the call.
-
-    Returns:
-      Published value.
-    '''
-    if len(self._devices) == 1:
-      return fn()
-
-    if name is None:
-      name = ops.get_default_graph().unique_name('pubsub')
-    name = name.replace(':', '_').replace('/', '_')
-
-    with PubSub._lock:
-      if name not in PubSub._name_to_call:
-        call = PubSub._prev_call
-        PubSub._prev_call += 1
-        PubSub._name_to_call[name] = call
-      else:
-        call = PubSub._name_to_call[name]
-
-    with ops.name_scope(f'{name}/replica{self._rank}'):
-      if self._root_rank != self._rank:
-        with ops.device(self._bcast_device):
-          return collective_ops.broadcast_recv(
-            shape, dtype,
-            len(self._devices),
-            self.channel,
-            call)
-      value = fn()
-      with ops.device(self._bcast_device):
-        bcast_send = collective_ops.broadcast_send(
-          value, shape, dtype,
-          len(self._devices),
-          self.channel,
-          call)
-      with ops.control_dependencies([bcast_send]):
-        return array_ops.identity(value)
diff --git a/hybridbackend/tensorflow/distribute/rpc.py b/hybridbackend/tensorflow/distribute/rpc.py
new file mode 100644
index 00000000..75d1feb8
--- /dev/null
+++ b/hybridbackend/tensorflow/distribute/rpc.py
@@ -0,0 +1,147 @@
+# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+r'''Classes for RPC based collective communication.
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import threading
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import collective_ops
+
+
+class RpcCollective(object):  # pylint: disable=useless-object-inheritance
+  r'''RPC based collective ops.
+  '''
+  _lock = threading.Lock()
+
+  _prev_group_key = 1
+  _group_key_dict = {}
+
+  _prev_instance_key = 1000000
+  _instance_key_dict = {}
+
+  def __init__(self, group_world, group_rank):
+    r'''Constructs an RPC based collective instance.
+
+    Args:
+      group_world: Involved devices.
+      group_rank: Rank of this peer.
+    '''
+    self._group_devices = sorted(group_world)
+    group_name = ','.join(self._group_devices)
+    with RpcCollective._lock:
+      if group_name not in RpcCollective._group_key_dict:
+        group_key = RpcCollective._prev_group_key
+        RpcCollective._prev_group_key += 1
+        RpcCollective._group_key_dict[group_name] = group_key
+    self._group_key = RpcCollective._group_key_dict[group_name]
+    self._group_rank = group_rank
+
+  @property
+  def group_key(self):
+    r'''Key for involved devices.
+    '''
+    return self._group_key
+
+  @property
+  def group_size(self):
+    r'''Size of involved devices.
+    '''
+    return len(self._group_devices)
+
+  @contextlib.contextmanager
+  def instance_scope(self, shared_name, name=None):
+    r'''Context for an RPC based collective op instance.
+    '''
+    if name is None:
+      name = ops.get_default_graph().unique_name(f'rpc_{shared_name}')
+    name = name.replace(':', '_').replace('/', '_')
+    with RpcCollective._lock:
+      if name not in RpcCollective._instance_key_dict:
+        instance_key = RpcCollective._prev_instance_key
+        RpcCollective._prev_instance_key += 1
+        RpcCollective._instance_key_dict[name] = instance_key
+      else:
+        instance_key = RpcCollective._instance_key_dict[name]
+    with ops.name_scope(f'{name}/replicas/{self._group_rank}/'):
+      yield instance_key
+
+  def broadcast(self, fn, dtype, shape, root_rank=None, name=None):
+    r'''Broadcast tensor across devices.
+
+    Args:
+      fn: Function to generate tensor.
+      dtype: Data type of the generated tensor.
+      shape: Shape of the generate tensor.
+      root_rank: Rank of the root peer.
+      name: Name of the op.
+
+    Returns:
+      Broadcasted value.
+    '''
+    if self.group_size == 1:
+      return fn()
+
+    if root_rank is None:
+      root_rank = 0
+    with self.instance_scope('broadcast', name=name) as instance_key:
+      if root_rank != self._group_rank:
+        with ops.device('/cpu:0'):
+          return collective_ops.broadcast_recv(
+            shape, dtype,
+            self.group_size,
+            self.group_key,
+            instance_key,
+            communication_hint='ring')
+      value = fn()
+      with ops.device('/cpu:0'):
+        bcast_send = collective_ops.broadcast_send(
+          value, shape, dtype,
+          self.group_size,
+          self.group_key,
+          instance_key,
+          communication_hint='ring')
+      with ops.control_dependencies([bcast_send]):
+        return array_ops.identity(value)
+
+  def allreduce(self, value, merge_op='Add', final_op='Id', name=None):
+    r'''Allreduce tensor across devices.
+
+    Args:
+      value: tensor to allreduce.
+      merge_op: 'Mul', 'Add'.
+      final_op: 'Id', 'Div'.
+      name: Name of the op.
+
+    Returns:
+      Allreduced value.
+    '''
+    if self.group_size == 1:
+      return value
+
+    with self.instance_scope('allreduce', name=name) as instance_key:
+      with ops.device('/cpu:0'):
+        return collective_ops.all_reduce(
+          value, self.group_size, self.group_key, instance_key,
+          merge_op=merge_op,
+          final_op=final_op,
+          communication_hint='ring')
diff --git a/hybridbackend/tensorflow/estimator/estimator.py b/hybridbackend/tensorflow/estimator/estimator.py
index f2b560a7..efc5266c 100644
--- a/hybridbackend/tensorflow/estimator/estimator.py
+++ b/hybridbackend/tensorflow/estimator/estimator.py
@@ -24,8 +24,6 @@
 
 from tensorflow.python.distribute import estimator_training
 from tensorflow.python.eager import context as _context
-from tensorflow.python.estimator import estimator as _estimator
-from tensorflow.python.estimator import run_config
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.ops import variable_scope as vs
@@ -37,20 +35,25 @@
 
 try:
   from tensorflow_estimator.python.estimator import estimator as _estimator_lib
+  from tensorflow_estimator.python.estimator import \
+    run_config as run_config_lib
   from tensorflow_estimator.python.estimator.export import export_lib
+  from tensorflow_estimator.python.estimator.training import _assert_eval_spec
   from tensorflow_estimator.python.estimator.training import _TrainingExecutor
 except ImportError:
   # pylint: disable=ungrouped-imports
   from tensorflow.python.estimator import estimator as _estimator_lib  # pylint: disable=reimported
+  from tensorflow.python.estimator import run_config as run_config_lib
   from tensorflow.python.estimator.export import export_lib
   from tensorflow.python.estimator.training import _TrainingExecutor
+  from tensorflow.python.estimator.training import _assert_eval_spec
   # pylint: enable=ungrouped-imports
 
+from hybridbackend.tensorflow.framework.config import wraps_session_config
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.device import device_function
 from hybridbackend.tensorflow.framework.ops import ModeKeys
 from hybridbackend.tensorflow.framework.rewriting import scope
-from hybridbackend.tensorflow.training.config import configure
 from hybridbackend.tensorflow.training.evaluation import EvaluationHook
 from hybridbackend.tensorflow.training.evaluation import EvaluationSpec
 from hybridbackend.tensorflow.training.hooks import Policy
@@ -62,7 +65,7 @@
 from hybridbackend.tensorflow.training.variables import reuse_variables
 
 
-class RunConfig(run_config.RunConfig):
+class RunConfig(run_config_lib.RunConfig):
   r'''RunConfig for estimators.
   '''
   @classmethod
@@ -73,7 +76,7 @@ def build(cls, prototype=None, **kwargs):
       return cls(**kwargs)
     prototype = prototype.replace(device_fn=device_function)
     prototype._is_chief = True  # pylint: disable=protected-access
-    prototype._session_config = configure(prototype=prototype.session_config)  # pylint: disable=protected-access
+    prototype._session_config = wraps_session_config(prototype.session_config)  # pylint: disable=protected-access
     if prototype._evaluation_master == '':  # pylint: disable=protected-access
       prototype._evaluation_master = prototype.master  # pylint: disable=protected-access
 
@@ -82,8 +85,8 @@ def build(cls, prototype=None, **kwargs):
   def __init__(self, **kwargs):
     r'''Creates a wrapped RunConfig.
     '''
-    kwargs['session_config'] = configure(
-      prototype=kwargs.pop('session_config', None))
+    kwargs['session_config'] = wraps_session_config(
+      kwargs.pop('session_config', None))
     kwargs['device_fn'] = device_function
     super().__init__(**kwargs)
     self._is_chief = True  # pylint: disable=protected-access
@@ -200,7 +203,7 @@ def _actual_eval(
           comm_pool_name=ModeKeys.EVAL,
           mode=ModeKeys.EVAL,
           model_dir=self._model_dir):
-        hooks = _estimator._check_hooks_type(hooks)  # pylint: disable=protected-access
+        hooks = _estimator_lib._check_hooks_type(hooks)  # pylint: disable=protected-access
         hooks.extend(self._convert_eval_steps_to_hooks(steps))  # pylint: disable=protected-access
         if not checkpoint_path:
           latest_path = checkpoint_management.latest_checkpoint(self._model_dir)  # pylint: disable=protected-access
@@ -254,13 +257,13 @@ def _eval_fn():
               hooks=hooks,
               update_op=update_op,
               eval_dict=eval_dict)
-        with Context.scope(eval_dir=self.eval_dir()):
-          eval_hook = EvaluationHook(
-            _eval_fn,
-            steps=eval_spec.steps,
-            every_n_iter=eval_every_n_iter,
-            history=eval_history)
-          train_hooks.append(eval_hook)
+        eval_hook = EvaluationHook(
+          _eval_fn,
+          steps=eval_spec.steps,
+          every_n_iter=eval_every_n_iter,
+          summary_dir=self.eval_dir(),
+          history=eval_history)
+        train_hooks.append(eval_hook)
 
       if self.config.cluster_spec:
         executor = TrainingExecutor(
@@ -272,14 +275,12 @@ def _eval_fn():
           raise ValueError(
             'Running `train_and_evaluate` with Distribute Coordinator '
             'not supported.')
-        with Context.scope(eval_dir=self.eval_dir()):
-          return executor.run()
+        return executor.run()
 
-      with Context.scope(eval_dir=self.eval_dir()):
-        return self.train(
-          train_spec.input_fn,
-          hooks=tuple(train_spec.hooks) + tuple(train_hooks),
-          max_steps=train_spec.max_steps)
+      return self.train(
+        train_spec.input_fn,
+        hooks=tuple(train_spec.hooks) + tuple(train_hooks),
+        max_steps=train_spec.max_steps)
 
     def export_saved_model(
         self, export_dir_base, serving_input_receiver_fn,
@@ -403,7 +404,6 @@ def predict(
       with _context.graph_mode(), Context.scope(
           mode=ModeKeys.PREDICT,
           model_dir=self._model_dir,
-          eval_dir=self.eval_dir(),
           comm_pool_capacity=1,
           comm_pool_name=ModeKeys.PREDICT):
         hooks = _estimator_lib._check_hooks_type(hooks)  # pylint: disable=protected-access
@@ -456,7 +456,7 @@ def predict(
   return HybridBackendEstimator
 
 
-Estimator = wraps_estimator(_estimator.Estimator)
+Estimator = wraps_estimator(_estimator_lib.Estimator)
 
 
 class TrainingExecutor(_TrainingExecutor):
@@ -469,27 +469,37 @@ def _start_std_server(self, config):
       config.cluster_spec,
       job_name=config.task_type,
       task_index=config.task_id,
-      config=configure(prototype=config.session_config),
+      config=wraps_session_config(config.session_config),
       start=True,
       protocol=config.protocol)
 
 
-def train_and_evaluate(
-    estimator, train_spec, eval_spec,
-    eval_every_n_iter=None,
-    eval_history=None):
+def train_and_evaluate(estimator, train_spec, eval_spec, **kwargs):
   r'''Train and evaluate the `estimator`.
-
-  Args:
-    eval_every_n_iter: `int`, runs parallel evaluation once every
-      N training iteration. If None, disable the evaluation.
-    eval_history: History of eval metrics. eval_history should support `append`
-      method.
   '''
-  if not isinstance(estimator, Estimator):
-    raise TypeError('estimator must be `hb.estimator.Estimator`')
-
-  return estimator.train_and_evaluate(
-    train_spec, eval_spec,
-    eval_every_n_iter=eval_every_n_iter,
-    eval_history=eval_history)
+  if hasattr(estimator, 'train_and_evaluate'):
+    return estimator.train_and_evaluate(train_spec, eval_spec, **kwargs)
+
+  _assert_eval_spec(eval_spec)  # fail fast if eval_spec is invalid.
+
+  executor = TrainingExecutor(
+    estimator=estimator,
+    train_spec=train_spec,
+    eval_spec=eval_spec)
+  config = estimator.config
+
+  # If `distribute_coordinator_mode` is set and running in distributed
+  # environment, we run `train_and_evaluate` via distribute coordinator.
+  if estimator_training.should_run_distribute_coordinator(config):
+    logging.info('Running `train_and_evaluate` with Distribute Coordinator.')
+    estimator_training.train_and_evaluate(
+      estimator, train_spec, eval_spec, TrainingExecutor)
+    return None
+
+  if (config.task_type == run_config_lib.TaskType.EVALUATOR
+      and config.task_id > 0):
+    raise ValueError(
+      'For distributed training, there can only be one `evaluator` task '
+      f'(with task id 0).  Given task id {config.task_id}')
+
+  return executor.run()
diff --git a/hybridbackend/tensorflow/feature_column/__init__.py b/hybridbackend/tensorflow/feature_column/__init__.py
deleted file mode 100644
index 53e1410e..00000000
--- a/hybridbackend/tensorflow/feature_column/__init__.py
+++ /dev/null
@@ -1,58 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Feature columns for fully sharded training.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import sys as _sys
-
-from tensorflow.python.framework import dtypes as _dtypes
-
-from hybridbackend.tensorflow.feature_column.backend import \
-  EmbeddingBackend as Backend
-from hybridbackend.tensorflow.feature_column.default import \
-  EmbeddingBackendDefault
-from hybridbackend.tensorflow.feature_column.dense_features import \
-  DenseFeatures
-from hybridbackend.tensorflow.feature_column.feature_column import \
-  raw_categorical_column
-from hybridbackend.tensorflow.framework.context import Context as _ctx
-from hybridbackend.tensorflow.framework.options import DictOption as _dict
-
-_ = (
-  _ctx.get().options
-  .register('emb_backend', 'DEFAULT', env='HB_EMB_BACKEND')
-  .register('emb_wire_dtype', _dict(_dtypes.float32), env='HB_EMB_WIRE_DTYPE')
-  .register('emb_buffer_size', 0, env='HB_EMB_BUFFER_SIZE')
-  .register('emb_buffer_load_factor', 0.5, env='HB_EMB_BUFFER_LOAD_FACTOR')
-  .register('emb_num_groups', 1, env='HB_EMB_NUM_GROUPS')
-  .register('emb_enable_concat', True)
-  .register('emb_num_buckets', _dict(0))
-  .register('emb_num_buckets_max', _dict(_sys.maxsize))
-  .register('emb_dimension', _dict(0))
-  .register('emb_combiner', _dict('mean'))
-  .register('emb_sharded', _dict(True))
-  .register('emb_unique', _dict(False))
-  .register('emb_pad', _dict(False))
-  .register('emb_device', _dict(''))
-  .register('emb_input_device', _dict(''))
-  .register('emb_dtype', _dict(_dtypes.float32))
-  .register('emb_input_dtype', _dict(_dtypes.int64))
-  .register('emb_segment_rank', _dict(0))
-  .register('emb_segment_sort', False, env='HB_EMB_SEGMENT_SORT'))
diff --git a/hybridbackend/tensorflow/feature_column/backend.py b/hybridbackend/tensorflow/feature_column/backend.py
deleted file mode 100644
index a6ed5576..00000000
--- a/hybridbackend/tensorflow/feature_column/backend.py
+++ /dev/null
@@ -1,236 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Backend of embedding tables.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import abc
-
-from hybridbackend.tensorflow.framework.context import Context
-from hybridbackend.tensorflow.framework.ops import ModeKeys
-
-
-class EmbeddingBackend(object):  # pylint: disable=useless-object-inheritance
-  r'''Embedding lookup backend.
-
-  An embedding backend manages underlying storage for embedding lookups. Data
-  scientists can extend this class for customized embedding layers.
-  '''
-  _registry = {}
-
-  @classmethod
-  def register(cls, impl):
-    r'''Register implementation.
-
-    Args:
-      impl: Implementation to register.
-    '''
-    cls._registry[impl.NAME] = impl
-
-  @classmethod
-  def get(cls):
-    r'''Get an instance of registered implementation.
-
-    Returns:
-      An instance of registered implementation.
-    '''
-    backend = Context.get().options.emb_backend
-    if backend not in cls._registry:
-      raise ValueError(f'emb_backend is invalid: {backend}')
-    return cls._registry[backend]
-
-  @abc.abstractmethod
-  def build(
-      self,
-      column,
-      name,
-      shape,
-      dtype=None,
-      trainable=True,
-      use_resource=True,
-      initializer=None,
-      collections=None,
-      layer=None):
-    r'''Creates the embedding lookup weight.
-
-    Args:
-      column: Name of the input column.
-      name: Name of the embedding weights.
-      shape: Shape of the embedding weights.
-      dtype: (Optional.) Data type of the embedding weights.
-      trainable: (Optional.) If True, the embedding weights can be trained.
-      use_resource: (Optional.) If True, the embedding weights uses resource.
-      initializer: (Optional.) Initializer of the embedding weights.
-      collections: (Optional.) Collections of the embedding weights.
-      layer: `DenseFeatures` layer which manages weights.
-
-    Returns:
-      The embedding weight.
-    '''
-
-  @abc.abstractmethod
-  def init_from_checkpoint(
-      self, column, ckpt_dir_or_file, tensor_name_in_ckpt, to_restore):
-    r'''Replaces initializers of embedding weights to load from checkpoints.
-
-    Args:
-      column: Name of the input column.
-      ckpt_dir_or_file: Directory with checkpoints file or path to checkpoint
-        or checkpoint file path.
-      tensor_name_in_ckpt: Name of the `Tensor` in `ckpt_dir_or_file` from
-        which  to restore the embedding weights. Required if `ckpt_to_load_from`
-        is not `None`.
-      to_restore: `Tensor` to restore.
-    '''
-
-  @abc.abstractmethod
-  def lookup(self, column, weight, inputs, sharded=False, buffered=False):
-    r'''Lookup for embedding vectors.
-
-    Args:
-      column: Name of the input column.
-      weight: Embedding weight.
-      inputs: Inputs for embedding lookup.
-      sharded: If True inputs are sharded.
-      buffered: If True initialization would be delayed.
-
-    Returns:
-      Embedding vectors from the weight.
-    '''
-
-  @abc.abstractmethod
-  def update(self, column, weight, indexed_updates):
-    r'''Update embedding weight.
-
-    Args:
-      column: Name of the input column.
-      weight: Embedding weight.
-      indexed_updates: An `IndexedSlices` to update weight in specific indices.
-    '''
-
-  @property
-  def buffer_size(self):
-    r'''Buffer size of embedding variables for training.
-    '''
-    return Context.get().options.emb_buffer_size
-
-  @property
-  def buffer_load_factor(self):
-    r'''Buffer load factor of embedding variables for training.
-    '''
-    return Context.get().options.emb_buffer_load_factor
-
-  @property
-  def num_groups(self):
-    r'''Number of embedding lookup groups for training.
-    '''
-    return Context.get().options.emb_num_groups
-
-  @property
-  def enable_concat(self):
-    r'''If True, concat embedding vectors for training.
-    '''
-    return Context.get().options.emb_enable_concat
-
-  @property
-  def enable_sharding(self):
-    r'''Whether the sharding is enabled.
-    '''
-    return Context.get().options.mode != ModeKeys.PREDICT
-
-  def num_buckets(self, column):
-    r'''Number of buckets for the column.
-    '''
-    return Context.get().options.emb_num_buckets[column]
-
-  def dimension(self, column):
-    r'''Dimension of the column.
-    '''
-    return Context.get().options.emb_dimension[column]
-
-  def combiner(self, column):
-    r'''Combiner of the column.
-    '''
-    return Context.get().options.emb_combiner[column]
-
-  def sharded(self, column):
-    r'''Whether the column should be sharded.
-    '''
-    if Context.get().world_size <= 1:
-      return False
-    if not self.enable_sharding:
-      return False
-    return Context.get().options.emb_sharded[column]
-
-  def unique(self, column):
-    r'''Whether inputs for the column is already unique.
-    '''
-    return Context.get().options.emb_unique[column]
-
-  def pad(self, column):
-    r'''Whether lookup results of the column should be padded.
-    '''
-    return Context.get().options.emb_pad[column]
-
-  def device(self, column):
-    r'''Device of the embedding weights.
-    '''
-    return Context.get().options.emb_device[column]
-
-  def input_device(self, column):
-    r'''Device of embedding lookup inputs.
-    '''
-    return Context.get().options.emb_input_device[column]
-
-  def dtype(self, column):
-    r'''Data type of the embedding weights.
-    '''
-    return Context.get().options.emb_dtype[column]
-
-  def wire_dtype(self, column):
-    r'''Data type of the column for communicaiton.
-    '''
-    return Context.get().options.emb_wire_dtype[column]
-
-  def input_dtype(self, column):
-    r'''Data type of the column inputs.
-    '''
-    return Context.get().options.emb_input_dtype[column]
-
-  def segment_rank(self, column):
-    r'''Segment rank of the column weights.
-    '''
-    return Context.get().options.emb_segment_rank[column]
-
-  def weight_name(self, column):
-    r'''Name of the column weights.
-    '''
-    name = 'embedding_weights'
-    if self.sharded(column):
-      shard = Context.get().rank
-      name = f'{name}/part_{shard}'
-    return name
-
-  def weight_shared_name(self, column, var):
-    r'''Get shared name of the column weights from an variable.
-    '''
-    var_name = var.name.split(':')[0]
-    if self.sharded(column):
-      return var_name.split('/part')[0]
-    return var_name
diff --git a/hybridbackend/tensorflow/feature_column/benchmark.py b/hybridbackend/tensorflow/feature_column/benchmark.py
deleted file mode 100644
index ab0bd62c..00000000
--- a/hybridbackend/tensorflow/feature_column/benchmark.py
+++ /dev/null
@@ -1,167 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Recmendation model training benchmark.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import argparse
-import math
-import re
-
-import tensorflow as tf
-
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-def build_bench_op(params):
-  with tf.device('/cpu:0'):
-    ds = hb.data.ParquetDataset(
-      params.filenames,
-      batch_size=params.batch_size,
-      num_parallel_reads=len(params.filenames),
-      drop_remainder=True)
-    ds = ds.apply(hb.data.parse())
-
-    def map_fn(batch):
-      labels = {}
-      numerics = {}
-      ids = {}
-      for f in batch:
-        if re.match(params.label_field_pattern, f):
-          labels[f] = tf.to_float(batch[f])
-        elif re.match(params.numeric_field_pattern, f):
-          numerics[f] = tf.to_float(batch[f])
-        elif re.match(params.categorical_field_pattern, f):
-          sp_ids = batch[f]
-          if isinstance(sp_ids, tf.Tensor):
-            sp_ids = tf.sparse.from_dense(
-              tf.reshape(sp_ids, (params.batch_size, 1)))
-          ids[f] = tf.SparseTensor(
-            sp_ids.indices,
-            sp_ids.values % params.bucket_size,
-            sp_ids.dense_shape)
-      return labels, numerics, ids
-
-    ds = ds.map(map_fn)
-    ds = ds.prefetch(1)
-    iterator = tf.data.make_one_shot_iterator(ds)
-  labels, numerics, ids = iterator.get_next()
-  embedding_weights = {}
-  embeddings = {}
-  for f, v in ids.items():
-    embedding_weight_name = f'{f}_weight'
-    with tf.device('/cpu:0'):
-      embedding_weights[embedding_weight_name] = tf.get_variable(
-        f'{f}_weight',
-        shape=(params.bucket_size, params.dimension),
-        initializer=tf.random_uniform_initializer(-1e-3, 1e-3))
-    embeddings[f] = tf.nn.safe_embedding_lookup_sparse(
-      embedding_weights[embedding_weight_name], v,
-      default_id=params.default_id)
-  labels = tf.concat(
-    [tf.reshape(v, [-1, 1]) for v in labels.values()],
-    axis=1)
-  numerics = tf.concat(
-    [tf.reshape(v, [-1, 1]) for v in numerics.values()],
-    axis=1)
-  embeddings = tf.concat(
-    [tf.reshape(v, [-1, params.dimension]) for v in embeddings.values()],
-    axis=1)
-  mlp_dims = [int(d) for d in params.mlp_dims.split(',') if d]
-  mlp_input = tf.concat([numerics, embeddings], axis=1)
-
-  for i, d in enumerate(mlp_dims):
-    mlp_input = tf.layers.dense(
-      mlp_input, d,
-      activation=tf.nn.relu,
-      kernel_initializer=tf.random_normal_initializer(
-        mean=0.0,
-        stddev=math.sqrt(2.0 / d)),
-      bias_initializer=tf.random_normal_initializer(
-        mean=0.0,
-        stddev=math.sqrt(1.0 / d)),
-      name=f'mlp_{i}')
-  logits = tf.layers.dense(
-    mlp_input, 1,
-    activation=tf.nn.sigmoid,
-    kernel_initializer=tf.random_normal_initializer(
-      mean=0.0,
-      stddev=math.sqrt(2.0)),
-    bias_initializer=tf.random_normal_initializer(
-      mean=0.0,
-      stddev=math.sqrt(1.0)),
-    name='logits')
-
-  loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(labels, logits))
-  mlp_weights = {v.name: v for v in tf.trainable_variables()}
-  for _, v in embedding_weights.items():
-    del mlp_weights[v.name]
-
-  step = tf.train.get_or_create_global_step()
-
-  embedding_opt = tf.train.AdagradOptimizer(learning_rate=params.lr)
-  embedding_train_op = embedding_opt.minimize(
-    loss, global_step=step, var_list=list(embedding_weights.values()))
-
-  mlp_opt = tf.train.AdamOptimizer(learning_rate=params.lr)
-  mlp_train_op = mlp_opt.minimize(
-    loss, global_step=step, var_list=list(mlp_weights.values()))
-
-  return tf.group([embedding_train_op, mlp_train_op])
-
-
-def benchmark(params):
-  with tf.device('/gpu:0'):
-    bench_op = build_bench_op(params)
-  hooks = [tf.train.StopAtStepHook(params.num_steps)]
-  hooks.append(
-    hb.train.StepStatHook(count=params.batch_size, unit='sample'))
-  if params.profile:
-    hooks.append(
-      tf.train.ProfilerHook(
-        save_steps=params.num_steps - 1,
-        output_dir='.'))
-  config = tf.ConfigProto(allow_soft_placement=True)
-  config.gpu_options.allow_growth = True
-  config.gpu_options.force_gpu_compatible = True
-  with tf.train.MonitoredTrainingSession(
-      '', hooks=hooks, config=config) as sess:
-    while not sess.should_stop():
-      sess.run(bench_op)
-
-
-if __name__ == '__main__':
-  hb.enable_optimization(logging_level=2)
-  tf.logging.set_verbosity(tf.logging.INFO)
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--lr', type=float, default=1.)
-  parser.add_argument('--dimension', type=int, default=32)
-  parser.add_argument('--bucket-size', type=int, default=100000)
-  parser.add_argument('--batch-size', type=int, default=100000)
-  parser.add_argument('--num-steps', type=int, default=100)
-  parser.add_argument('--profile', default=False, action='store_true')
-  parser.add_argument('--label-field-pattern', default='^label$')
-  parser.add_argument('--categorical-field-pattern', default='^id.*')
-  parser.add_argument('--numeric-field-pattern', default='^f.*')
-  parser.add_argument('--categorical-field-missing-value', type=int, default=0)
-  parser.add_argument('--default-id', type=int, default=None)
-  parser.add_argument('--mlp-dims', default='')
-  parser.add_argument('filenames', nargs='+')
-  benchmark(parser.parse_args())
diff --git a/hybridbackend/tensorflow/feature_column/default.py b/hybridbackend/tensorflow/feature_column/default.py
deleted file mode 100644
index ab44a464..00000000
--- a/hybridbackend/tensorflow/feature_column/default.py
+++ /dev/null
@@ -1,170 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Default backend of embedding tables.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.framework import dtypes
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import variable_scope as vs
-from tensorflow.python.ops import variables
-from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.training import checkpoint_utils
-
-from hybridbackend.tensorflow.feature_column.backend import EmbeddingBackend
-from hybridbackend.tensorflow.framework.context import Context
-
-
-class EmbeddingBackendDefault(EmbeddingBackend):
-  r'''Embedding backend for variables.
-  '''
-  NAME = 'DEFAULT'
-
-  def _get_variable(
-      self,
-      name,
-      shape=None,
-      dtype=dtypes.float32,
-      initializer=None,
-      trainable=None,
-      constraint=None,
-      use_resource=None,
-      collections=None,
-      synchronization=variables.VariableSynchronization.AUTO,
-      aggregation=variables.VariableAggregation.NONE,
-      partitioner=None):  # pylint: disable=unused-argument
-    r'''Gets an existing variable with these parameters or create a new one.
-    '''
-    return vs.get_variable(
-      name,
-      shape=shape,
-      dtype=dtype,
-      initializer=initializer,
-      trainable=trainable,
-      constraint=constraint,
-      use_resource=use_resource,
-      collections=collections,
-      synchronization=synchronization,
-      aggregation=aggregation,
-      partitioner=partitioner)
-
-  def build(
-      self,
-      column,
-      name,
-      shape,
-      dtype=None,
-      trainable=True,
-      use_resource=True,
-      initializer=None,
-      collections=None,
-      layer=None):
-    r'''Creates the embedding lookup weight.
-    '''
-    del use_resource
-    num_buckets_max = Context.get().options.emb_num_buckets_max[column]
-    num_buckets, dim = shape
-    if num_buckets is None or num_buckets >= num_buckets_max:
-      raise ValueError(
-        f'num_buckets of column {column} must be set explicitly')
-    if self.sharded(column):
-      shard = Context.get().rank
-      num_shards = Context.get().world_size
-      sharded_bucket_size = num_buckets // num_shards
-      if shard < num_buckets % num_shards:
-        sharded_bucket_size += 1
-      shape = (sharded_bucket_size, dim)
-    if layer is None:
-      var = vs.get_variable(
-        name=name,
-        shape=shape,
-        dtype=dtype,
-        initializer=initializer,
-        trainable=trainable,
-        use_resource=False,
-        collections=collections)
-    else:
-      var = layer.add_variable(
-        name=name,
-        shape=shape,
-        dtype=dtype,
-        initializer=initializer,
-        trainable=trainable,
-        use_resource=False,
-        getter=self._get_variable,
-        collections=collections)
-    if self.sharded(column):
-      num_shards = Context.get().world_size
-      shard = Context.get().rank
-      bucket_offset = (num_buckets // num_shards) * shard
-      remained_buckets = num_buckets % num_shards
-      if shard < remained_buckets:
-        bucket_offset += shard
-      else:
-        bucket_offset += remained_buckets
-      full_name = self.weight_shared_name(column, var)
-      if hasattr(var, '_set_save_slice_info'):
-        var._set_save_slice_info(  # pylint: disable=protected-access
-          variables.Variable.SaveSliceInfo(
-            full_name=full_name,
-            full_shape=[num_buckets, dim],
-            var_offset=[bucket_offset, 0],
-            var_shape=shape))
-      elif isinstance(var, variables.PartitionedVariable):
-        for pvar in var:
-          poffset = pvar._get_save_slice_info().var_offset  # pylint: disable=protected-access
-          pvar._set_save_slice_info(  # pylint: disable=protected-access
-            variables.Variable.SaveSliceInfo(
-              full_name=full_name,
-              full_shape=[num_buckets, dim],
-              var_offset=[bucket_offset + poffset[0], poffset[1]],
-              var_shape=pvar.shape))
-      else:
-        logging.warning(f'Saving variable without elasticity: {var}')
-      return var
-    return var
-
-  def init_from_checkpoint(
-      self, column, ckpt_dir_or_file, tensor_name_in_ckpt, to_restore):
-    r'''Replaces initializers of embedding weights to load from checkpoints.
-    '''
-    del column
-    if isinstance(to_restore, variables.PartitionedVariable):
-      to_restore = to_restore._get_variable_list()  # pylint: disable=protected-access
-    checkpoint_utils.init_from_checkpoint(
-      ckpt_dir_or_file, {tensor_name_in_ckpt: to_restore})
-
-  def lookup(self, column, weight, inputs, sharded=False, buffered=False):
-    r'''Lookup for embedding vectors.
-    '''
-    del column
-    del buffered
-    if sharded:
-      num_shards = Context.get().world_size
-      return array_ops.gather(weight, inputs // num_shards)
-    return array_ops.gather(weight, inputs)
-
-  def update(self, column, weight, indexed_updates):
-    r'''Update embedding weight.
-    '''
-    del column
-    return weight.scatter_update(indexed_updates)
-
-
-EmbeddingBackend.register(EmbeddingBackendDefault())
diff --git a/hybridbackend/tensorflow/feature_column/dense_features.py b/hybridbackend/tensorflow/feature_column/dense_features.py
deleted file mode 100644
index a50234c8..00000000
--- a/hybridbackend/tensorflow/feature_column/dense_features.py
+++ /dev/null
@@ -1,224 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''A layer that produces a dense `Tensor` based on given `feature_columns`.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.feature_column import feature_column as fc_old
-from tensorflow.python.feature_column import feature_column_v2 as fc
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import array_ops
-
-from hybridbackend.tensorflow.feature_column.backend import EmbeddingBackend
-from hybridbackend.tensorflow.feature_column.feature_column import \
-  _get_sparse_tensors
-from hybridbackend.tensorflow.feature_column.feature_column import \
-  EmbeddingColumn
-from hybridbackend.tensorflow.feature_column.feature_column import \
-  SharedEmbeddingColumn
-from hybridbackend.tensorflow.feature_column.feature_column import \
-  StateManagerImpl
-from hybridbackend.tensorflow.feature_column.lookup import \
-  EmbeddingLookupCoalesced
-from hybridbackend.tensorflow.framework.context import Context
-
-# pylint: disable=ungrouped-imports
-try:
-  from tensorflow.python.feature_column.dense_features import \
-    DenseFeatures as _DenseFeatures
-except ImportError:
-  from tensorflow.python.feature_column.feature_column_v2 import \
-    FeatureLayer as _DenseFeatures
-# pylint: enable=ungrouped-imports
-
-
-class DenseFeatures(_DenseFeatures):
-  r'''A layer that produces a dense `Tensor` based on given `feature_columns`.
-  '''
-  def __init__(self, feature_columns, trainable=True, name=None, **kwargs):
-    r'''Constructs a DenseFeatures layer.
-    '''
-    self._impl = EmbeddingBackend.get()
-    self._num_groups = kwargs.pop('num_groups', self._impl.num_groups)
-    self._enable_concat = kwargs.pop('enable_concat', self._impl.enable_concat)
-    verified_feature_columns = []
-    for f in feature_columns:
-      if isinstance(f, fc.EmbeddingColumn):
-        verified_feature_columns.append(EmbeddingColumn.build(f))
-      elif isinstance(f, fc_old._SharedEmbeddingColumn):
-        verified_feature_columns.append(SharedEmbeddingColumn.build(f))
-      elif isinstance(f, fc.SharedEmbeddingColumn):
-        raise ValueError(
-          f'{f} not supported: Use tf.feature_column.shared_embeding_columns')
-      else:
-        verified_feature_columns.append(f)
-    super().__init__(
-      feature_columns=verified_feature_columns,
-      trainable=trainable,
-      name=name,
-      **kwargs)
-    self._lookup_sparse_coalesced = EmbeddingLookupCoalesced()
-    self._state_manager = StateManagerImpl(self, self.trainable)  # pylint: disable=protected-access
-
-  @property
-  def num_groups(self):
-    r'''Number of embedding column groups.
-    '''
-    return self._num_groups
-
-  def call(self, features, cols_to_output_tensors=None):
-    r'''Returns a dense tensor corresponding to the `feature_columns`.
-
-    Args:
-      features: A mapping from key to tensors. `FeatureColumn`s look up via
-        these keys. For example `numeric_column('price')` will look at 'price'
-        key in this dict. Values can be a `SparseTensor` or a `Tensor` depends
-        on corresponding `FeatureColumn`.
-      cols_to_output_tensors: If not `None`, this will be filled with a dict
-        mapping feature columns to output tensors created.
-
-    Returns:
-      A `Tensor` which represents input layer of a model. Its shape
-      is (batch_size, first_layer_dimension) and its dtype is `float32`.
-      first_layer_dimension is determined based on given `feature_columns`.
-      If emb_enable_concat is disabled, `None` would be returned.
-
-    Raises:
-      ValueError: If arguments are not valid.
-    '''
-    if not isinstance(features, dict):
-      raise ValueError('We expected a dictionary here. Instead we got: ',
-                       features)
-
-    for k, v in features.items():
-      if (isinstance(v, ops.Tensor)
-          and v.shape.rank is not None
-          and v.shape.rank > 1):
-        raise TypeError(
-          f'Column {k} has a multi-rank tf.Tensor input, please uses '
-          'tf.sparse.SparseTensor instead')
-
-    transformation_cache = fc.FeatureTransformationCache(features)
-    indexed_coalesced_columns = []
-    indexed_non_coalesced_columns = []
-    for cid, c in enumerate(self._feature_columns):
-      if self._impl.sharded(c.categorical_column.name):
-        indexed_coalesced_columns.append(tuple([cid, c]))
-      else:
-        indexed_non_coalesced_columns.append(tuple([cid, c]))
-
-    num_groups = self.num_groups
-    if num_groups is None or num_groups > len(indexed_coalesced_columns):
-      num_groups = len(indexed_coalesced_columns)
-    num_groups = max(num_groups, 1)
-
-    non_coalesced_column_output_tensors = []
-    with ops.name_scope(self.name):
-      for cid, c in indexed_non_coalesced_columns:
-        with ops.name_scope(c.name):
-          tensor = c.get_dense_tensor(transformation_cache, self._state_manager)
-          if hasattr(self, '_process_dense_tensor'):
-            processed_tensors = self._process_dense_tensor(c, tensor)
-          else:
-            processed_tensors = tensor
-          if cols_to_output_tensors is not None:
-            cols_to_output_tensors[c] = processed_tensors
-          non_coalesced_column_output_tensors.append(
-            tuple([cid, processed_tensors]))
-
-    coalesced_column_output_tensors = []
-    group_columns = []
-    group_inputs = []
-    group_weights = []
-    for cid, c in indexed_coalesced_columns:
-      with ops.name_scope(c.name):
-        group_columns.append(c)
-        sparse_tensors = _get_sparse_tensors(
-          c.categorical_column, transformation_cache, self._state_manager)
-        group_inputs.append(sparse_tensors)
-        if hasattr(c, 'get_state'):
-          group_weights.append(c.get_state(self._state_manager))
-        else:
-          group_weights.append(
-            self._state_manager.get_variable(
-              c, self._impl.weight_name(c.categorical_column.name)))
-    if indexed_coalesced_columns:
-      _, coalesced_columns = zip(*indexed_coalesced_columns)
-      coalesced_column_names = [
-        c.categorical_column.name for c in coalesced_columns]
-      column_dimension = {
-        c.categorical_column.name: c.dimension
-        for c in coalesced_columns}
-      column_combiner = {
-        c.categorical_column.name: c.combiner
-        for c in coalesced_columns}
-      column_num_buckets = {
-        c.categorical_column.name: getattr(
-          c.categorical_column, 'num_buckets',
-          getattr(c.categorical_column, '_num_buckets', None))
-        for c in coalesced_columns}
-      with Context.scope(
-          emb_dimension=column_dimension,
-          emb_combiner=column_combiner,
-          emb_num_buckets=column_num_buckets):
-        group_outputs = self._lookup_sparse_coalesced(
-          group_weights, group_inputs, coalesced_column_names,
-          num_groups=num_groups)
-    for idx, cid_and_column in enumerate(indexed_coalesced_columns):
-      cid, c = cid_and_column
-      with ops.name_scope(c.name):
-        if hasattr(self, '_process_dense_tensor'):
-          processed_tensors = self._process_dense_tensor(c, group_outputs[idx])
-        else:
-          processed_tensors = group_outputs[idx]
-        if cols_to_output_tensors is not None:
-          cols_to_output_tensors[c] = processed_tensors
-        coalesced_column_output_tensors.append(
-          tuple([cid, processed_tensors]))
-
-    indexed_output_tensors = sorted(
-      coalesced_column_output_tensors + non_coalesced_column_output_tensors)
-    _, output_tensors = zip(*indexed_output_tensors)
-
-    if not self._enable_concat:
-      return output_tensors
-
-    if hasattr(self, '_verify_and_concat_tensors'):
-      return self._verify_and_concat_tensors(output_tensors)
-    return array_ops.concat(output_tensors, 1)
-
-
-def dense_features(features, feature_columns):
-  r'''Function produces dense tensors based on given `feature_columns`.
-
-  Args:
-    features: A mapping from key to tensors. `FeatureColumn`s look up via
-      these keys. For example `numeric_column('price')` will look at 'price'
-      key in this dict. Values can be a `SparseTensor` or a `Tensor` depends
-      on corresponding `FeatureColumn`.
-    feature_columns: List of feature columns.
-
-  Returns:
-    List of `Tensor`s which represents input layer of a model, which matches
-    order of columns in `feature_columns`.
-  '''
-  cols_to_output_tensors = {}
-  DenseFeatures(feature_columns)(
-    features, cols_to_output_tensors=cols_to_output_tensors)
-  return [cols_to_output_tensors[f] for f in feature_columns]
diff --git a/hybridbackend/tensorflow/feature_column/feature_column.py b/hybridbackend/tensorflow/feature_column/feature_column.py
deleted file mode 100644
index 41f9d369..00000000
--- a/hybridbackend/tensorflow/feature_column/feature_column.py
+++ /dev/null
@@ -1,433 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Feature columns for embedding lookup.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import sys
-import threading
-
-from tensorflow.python.feature_column import feature_column as fc_old
-from tensorflow.python.feature_column import feature_column_v2 as fc
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import check_ops
-
-from hybridbackend.tensorflow.feature_column.backend import EmbeddingBackend
-from hybridbackend.tensorflow.feature_column.lookup import EmbeddingLookup
-from hybridbackend.tensorflow.framework.context import Context
-from hybridbackend.tensorflow.framework.ops import GraphKeys
-
-
-class PatchEmbeddingColumn(object):  # pylint: disable=useless-object-inheritance
-  r'''Context manager that patches fc.EmbeddingColumn.
-  '''
-  _lock = threading.Lock()
-  _stack_depth = 0
-
-  def __enter__(self):
-    with PatchEmbeddingColumn._lock:
-      PatchEmbeddingColumn._stack_depth += 1
-      if PatchEmbeddingColumn._stack_depth <= 1:
-
-        def decorate_get_raw_feature(fn):  # pylint: disable=unused-argument
-          r'''w/o expanding dims when rank equals one.
-          '''
-          def decorated(col, key):
-            raw_feature = col._features[key]  # pylint: disable=protected-access
-            feature_tensor = \
-              sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(raw_feature)
-            rank = feature_tensor.get_shape().ndims
-            if rank is not None:
-              if rank == 0:
-                raise ValueError(
-                  f'Feature (key: {key}) cannot have rank 0.'
-                  f'Given: {feature_tensor}')
-              return feature_tensor
-            # Handle dynamic rank.
-            with ops.control_dependencies([
-                check_ops.assert_positive(
-                  array_ops.rank(feature_tensor),
-                  message=f'Feature (key: {key}) cannot have rank 0.'
-                          f'Given: {feature_tensor}')]):
-              return feature_tensor
-          return decorated
-        self._prev_get_raw_feature = \
-          fc.FeatureTransformationCache._get_raw_feature_as_tensor
-        fc.FeatureTransformationCache._get_raw_feature_as_tensor = \
-          decorate_get_raw_feature(self._prev_get_raw_feature)
-
-      return self
-
-  def __exit__(self, exc_type, exc_val, exc_tb):
-    with PatchEmbeddingColumn._lock:
-      if PatchEmbeddingColumn._stack_depth <= 1:
-        fc.FeatureTransformationCache._get_raw_feature_as_tensor = \
-          self._prev_get_raw_feature
-      PatchEmbeddingColumn._stack_depth -= 1
-
-
-def _get_sparse_tensors(
-    categorical_column, transformation_cache, state_manager):
-  r'''Returns tensor before doing the embedding lookup.
-  '''
-  if isinstance(categorical_column, fc.SequenceCategoricalColumn):
-    categorical_column_type = type(categorical_column)
-    raise ValueError(
-      f'categorical_column {categorical_column.name} '
-      'must not be of type SequenceCategoricalColumn. '
-      'Suggested fix A: If you wish to use DenseFeatures, use a '
-      'non-sequence categorical_column_with_*. '
-      'Suggested fix B: If you wish to create sequence input, use '
-      'SequenceFeatures instead of DenseFeatures. '
-      f'Given (type {categorical_column_type}): '
-      f'{categorical_column}')
-
-  if isinstance(categorical_column, fc.IdentityCategoricalColumn):
-    dynamic_bucketing = (
-      Context.get().options.emb_num_buckets_max[
-        categorical_column.name])
-    if categorical_column.number_buckets >= dynamic_bucketing:
-      with PatchEmbeddingColumn():
-        return fc.CategoricalColumn.IdWeightPair(
-          transformation_cache.get(categorical_column.key, state_manager), None)
-
-  return categorical_column.get_sparse_tensors(
-    transformation_cache, state_manager)
-
-
-def raw_categorical_column(key):
-  r'''A `CategoricalColumn` that returns unchecked values.
-
-  Args:
-    key: A unique string identifying the input feature. It is used as the
-      column name and the dictionary key for feature parsing configs, feature
-      `Tensor` objects, and feature columns.
-  '''
-  return fc.categorical_column_with_identity(key, num_buckets=sys.maxsize)
-
-
-class StateManagerImpl(fc._StateManagerImpl):  # pylint: disable=protected-access
-  r'''Manages the state of DenseFeatures.
-  '''
-  def create_variable(
-      self,
-      feature_column,
-      name,
-      shape,
-      dtype=None,
-      trainable=True,
-      use_resource=True,
-      initializer=None):
-    r'''Creates a new variable.
-    '''
-    if name in self._cols_to_vars_map[feature_column]:
-      raise ValueError('Variable already exists.')
-
-    impl = EmbeddingBackend.get()
-    collections = [ops.GraphKeys.GLOBAL_VARIABLES]
-    if impl.sharded(feature_column.categorical_column.name):
-      collections.append(GraphKeys.SHARDED_VARIABLES)
-    with Context.scope(sharding=False):
-      var = impl.build(
-        feature_column.categorical_column.name,
-        name,
-        shape,
-        dtype=dtype,
-        trainable=self._trainable and trainable,
-        use_resource=use_resource,
-        initializer=initializer,
-        collections=collections,
-        layer=self._layer)
-    self._cols_to_vars_map[feature_column][name] = var
-    return var
-
-
-class EmbeddingColumn(fc.EmbeddingColumn):
-  r'''Feature column for dense features.
-  '''
-  @classmethod
-  def build(cls, prototype):
-    r'''Build EmbeddingColumn from prototype.
-    '''
-    if isinstance(prototype, cls):
-      return prototype
-    if isinstance(prototype, fc.EmbeddingColumn):
-      prototype_kwargs = dict(prototype._asdict())  # pylint: disable=protected-access
-      return cls(**prototype_kwargs)
-    raise ValueError('EmbeddingColumn required')
-
-  def __new__(
-      cls,
-      categorical_column,
-      dimension,
-      combiner,
-      initializer,
-      ckpt_to_load_from,
-      tensor_name_in_ckpt,
-      max_norm,
-      trainable,
-      **kwargs):
-    r'''Constructor.
-    '''
-    self = super(EmbeddingColumn, cls).__new__(
-      cls,
-      categorical_column,
-      dimension,
-      combiner,
-      initializer,
-      ckpt_to_load_from,
-      tensor_name_in_ckpt,
-      max_norm,
-      trainable,
-      **kwargs)
-    # pylint: disable=protected-access
-    self._impl = EmbeddingBackend.get()
-    self._lookup_sparse = EmbeddingLookup()
-    # pylint: enable=protected-access
-    return self
-
-  def _get_dense_tensor_internal_helper(
-      self, sparse_tensors, embedding_weights):
-    r'''Private method that follows the signature of get_dense_tensor.
-    '''
-    if self.ckpt_to_load_from is not None:
-      self._impl.init_from_checkpoint(
-        self.categorical_column.name,
-        self.ckpt_to_load_from,
-        self.tensor_name_in_ckpt,
-        embedding_weights)
-
-    with Context.scope(
-        emb_dimension=self.dimension,
-        emb_combiner=self.combiner,
-        emb_num_buckets=self.num_buckets):
-      return self._lookup_sparse(
-        self.categorical_column.name,
-        embedding_weights,
-        sparse_tensors,
-        name=f'{self.name}_lookup')
-
-  def _get_dense_tensor_internal(self, sparse_tensors, state_manager):
-    r'''Private method that follows the signature of get_dense_tensor.
-    '''
-    embedding_weights = self.get_state(state_manager)
-    return self._get_dense_tensor_internal_helper(sparse_tensors,
-                                                  embedding_weights)
-
-  @property
-  def num_buckets(self):
-    r'''Number of buckets.
-    '''
-    return getattr(
-      self.categorical_column, 'num_buckets',
-      getattr(self.categorical_column, '_num_buckets', None))
-
-  def get_dense_tensor(self, transformation_cache, state_manager):
-    r'''Returns tensor after doing the embedding lookup.
-    '''
-    sparse_tensors = _get_sparse_tensors(
-      self.categorical_column, transformation_cache, state_manager)
-    return self._get_dense_tensor_internal(sparse_tensors, state_manager)
-
-  def output_shape(self, inputs):
-    r'''Tuple of column output shape.
-    '''
-    inputs_shape = array_ops.shape(inputs)[:-1]
-    num_elements = self.variable_shape.num_elements()
-    actual_inputs_shape = array_ops.concat(
-      [inputs_shape, array_ops.constant(num_elements, shape=[1])],
-      0, name='actual_inputs_shape')
-    return actual_inputs_shape
-
-  def get_state(self, state_manager):
-    r'''Get or create embedding weights.
-    '''
-    return state_manager.get_variable(
-      self, self._impl.weight_name(self.categorical_column.name))
-
-  def create_state(self, state_manager):
-    r'''Uses the `state_manager` to create state for the FeatureColumn.
-    '''
-    embedding_shape = (self.num_buckets, self.dimension)
-    embedding_name = self._impl.weight_name(self.categorical_column.name)
-    with ops.device(self._impl.device(self.categorical_column.name)):
-      return state_manager.create_variable(
-        self,
-        name=embedding_name,
-        shape=embedding_shape,
-        dtype=self._impl.dtype(self.categorical_column.name),
-        trainable=self.trainable,
-        use_resource=True,
-        initializer=self.initializer)
-
-
-class SharedEmbeddingColumn(fc_old._SharedEmbeddingColumn, fc.DenseColumn):  # pylint: disable=protected-access
-  r'''Feature column for dense features.
-  '''
-  @classmethod
-  def build(cls, prototype):
-    r'''Build SharedEmbeddingColumn from prototype.
-    '''
-    if isinstance(prototype, cls):
-      return prototype
-    if isinstance(prototype, fc_old._SharedEmbeddingColumn):  # pylint: disable=protected-access
-      prototype_kwargs = dict(prototype._asdict())  # pylint: disable=protected-access
-      return cls(**prototype_kwargs)
-    raise ValueError('SharedEmbeddingColumn required')
-
-  def __new__(
-      cls,
-      categorical_column,
-      dimension,
-      combiner,
-      initializer,
-      shared_embedding_collection_name,
-      ckpt_to_load_from,
-      tensor_name_in_ckpt,
-      max_norm,
-      trainable,
-      **kwargs):
-    r'''Constructor.
-    '''
-    self = super(SharedEmbeddingColumn, cls).__new__(
-      cls,
-      categorical_column,
-      dimension,
-      combiner,
-      initializer,
-      shared_embedding_collection_name,
-      ckpt_to_load_from,
-      tensor_name_in_ckpt,
-      max_norm,
-      trainable,
-      **kwargs)
-    # pylint: disable=protected-access
-    self._impl = EmbeddingBackend.get()
-    self._lookup_sparse = EmbeddingLookup()
-    # pylint: enable=protected-access
-    return self
-
-  def _get_dense_tensor_internal_helper(
-      self, sparse_tensors, embedding_weights):
-    r'''Private method that follows the signature of get_dense_tensor.
-    '''
-    if self.ckpt_to_load_from is not None:
-      self._impl.init_from_checkpoint(
-        self.categorical_column.name,
-        self.ckpt_to_load_from,
-        self.tensor_name_in_ckpt,
-        embedding_weights)
-
-    with Context.scope(
-        emb_dimension=self.dimension,
-        emb_combiner=self.combiner,
-        emb_num_buckets=self.num_buckets):
-      return self._lookup_sparse(
-        self.categorical_column.name,
-        embedding_weights,
-        sparse_tensors,
-        name=f'{self.name}_lookup')
-
-  def _get_dense_tensor_internal(self, sparse_tensors, state_manager):
-    r'''Private method that follows the signature of get_dense_tensor.
-    '''
-    embedding_weights = self.get_state(state_manager)
-    return self._get_dense_tensor_internal_helper(sparse_tensors,
-                                                  embedding_weights)
-
-  def get_dense_tensor(self, transformation_cache, state_manager):
-    r'''Returns tensor after doing the embedding lookup.
-    '''
-    sparse_tensors = _get_sparse_tensors(
-      self.categorical_column, transformation_cache, state_manager)
-    return self._get_dense_tensor_internal(sparse_tensors, state_manager)
-
-  @property
-  def num_buckets(self):
-    r'''Number of buckets.
-    '''
-    return getattr(
-      self.categorical_column, 'num_buckets',
-      getattr(self.categorical_column, '_num_buckets', None))
-
-  @property
-  def variable_shape(self):
-    r'''See `DenseColumn` base class.
-    '''
-    return tensor_shape.TensorShape([self.dimension])
-
-  def output_shape(self, inputs):
-    r'''Tuple of column output shape.
-    '''
-    inputs_shape = array_ops.shape(inputs)[:-1]
-    num_elements = self.variable_shape.num_elements()
-    actual_inputs_shape = array_ops.concat(
-      [inputs_shape, array_ops.constant(num_elements, shape=[1])],
-      0, name='actual_inputs_shape')
-    return actual_inputs_shape
-
-  def get_state(self, state_manager):
-    r'''Get or create embedding weights.
-    '''
-    del state_manager
-    shared_embedding_collection = ops.get_collection(
-      self.shared_embedding_collection_name)
-    if not shared_embedding_collection:
-      raise ValueError('create_state must be called before')
-    if len(shared_embedding_collection) > 1:
-      raise ValueError(
-        f'Collection {shared_embedding_collection} can only contain one '
-        'variable. Suggested fix A: Choose a unique name for this '
-        'collection. Suggested fix B: Do not add any variables to this '
-        'collection. The feature_column library already adds a variable '
-        'under the hood.')
-    return shared_embedding_collection[0]
-
-  def create_state(self, state_manager):
-    r'''Uses the `state_manager` to create state for the FeatureColumn.
-    '''
-    shared_embedding_collection = ops.get_collection(
-      self.shared_embedding_collection_name)
-    if shared_embedding_collection:
-      if len(shared_embedding_collection) > 1:
-        raise ValueError(
-          f'Collection {shared_embedding_collection} can only contain one '
-          'variable. Suggested fix A: Choose a unique name for this '
-          'collection. Suggested fix B: Do not add any variables to this '
-          'collection. The feature_column library already adds a variable '
-          'under the hood.')
-      return shared_embedding_collection[0]
-
-    embedding_shape = (self.num_buckets, self.dimension)
-    embedding_name = self._impl.weight_name(self.categorical_column.name)
-    with ops.device(self._impl.device(self.categorical_column.name)):
-      embedding_weights = state_manager.create_variable(
-        self,
-        name=embedding_name,
-        shape=embedding_shape,
-        dtype=self._impl.dtype(self.categorical_column.name),
-        trainable=self.trainable,
-        use_resource=True,
-        initializer=self.initializer)
-      ops.add_to_collection(self.shared_embedding_collection_name,
-                            embedding_weights)
-      return embedding_weights
diff --git a/hybridbackend/tensorflow/feature_column/lookup.py b/hybridbackend/tensorflow/feature_column/lookup.py
deleted file mode 100644
index a36b3294..00000000
--- a/hybridbackend/tensorflow/feature_column/lookup.py
+++ /dev/null
@@ -1,288 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Functors for embedding lookup.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import gradients
-from tensorflow.python.platform import tf_logging as logging
-
-from hybridbackend.tensorflow.distribute.communicator_pool import \
-  CommunicatorPool
-from hybridbackend.tensorflow.feature_column.backend import EmbeddingBackend
-from hybridbackend.tensorflow.feature_column.math_lib import segment_reduce
-from hybridbackend.tensorflow.framework.context import Context
-from hybridbackend.tensorflow.ops.floormod_shuffle.ops import floormod_shuffle
-from hybridbackend.tensorflow.ops.floormod_shuffle.ops import \
-  floormod_shuffle_n
-
-
-class EmbeddingLookup(object):  # pylint: disable=useless-object-inheritance
-  r'''Functor to lookup embeddings.
-  '''
-  def __init__(self):
-    self._impl = EmbeddingBackend.get()
-
-  def _exchange_ids(self, comm, inputs, inputs_deps):
-    r'''Communicator call to exchange IDs.
-    '''
-    with ops.control_dependencies(inputs_deps):
-      return comm.alltoallv(*inputs), None
-
-  def _exchange_embeddings(self, dimension):
-    r'''Communicator call to exchange embeddings.
-    '''
-    def _fn(comm, inputs, inputs_deps):
-      r'''Communicator call to exchange embeddings.
-      '''
-      if inputs_deps is None:
-        embs, _ = comm.alltoallv(*inputs, common_shape=[dimension])
-        return embs, None
-      with ops.control_dependencies(inputs_deps):
-        embs, _ = comm.alltoallv(*inputs, common_shape=[dimension])
-
-      def grad_fn(grads, grads_deps, roots):
-        r'''Gradient function.
-        '''
-        del roots
-        with ops.control_dependencies(grads_deps):
-          d_inputs = gradients.gradients(embs, inputs, grad_ys=grads)
-        return d_inputs, None
-      return embs, grad_fn
-    return _fn
-
-  def __call__(self, column, weights, inputs, name=None):
-    r'''Lookup embedding results for sparse tensors.
-    '''
-    with ops.name_scope(name):
-      sparse_ids = inputs.id_tensor
-      sparse_weights = inputs.weight_tensor
-      if isinstance(sparse_ids, sparse_tensor.SparseTensor):
-        ids = sparse_ids.values
-      else:
-        ids = sparse_ids
-      input_device = self._impl.input_device(column)
-      pad = self._impl.pad(column)
-      segment_rank = self._impl.segment_rank(column)
-      with ops.device(input_device):
-        ids = array_ops.reshape(ops.convert_to_tensor(ids), [-1])
-        if self._impl.unique(column):
-          unique_index = None
-          if not pad:
-            logging.info('If unique is set to True, pad would be set to True')
-            pad = True
-        else:
-          ids, unique_index = array_ops.unique(ids)
-        if segment_rank != 0 and not pad:
-          logging.info('If segment rank is not 0, pad would be set to True')
-          pad = True
-      dimension = self._impl.dimension(column)
-      if self._impl.sharded(column):
-        with ops.name_scope('shuffle_ids'):
-          ids_shards, ids_sizes, partition_index = floormod_shuffle(
-            ids, Context.get().world_size)
-          shard_ids, embs_sizes = CommunicatorPool.get().alltoallv(
-            ids_shards, ids_sizes,
-            trainable=False)
-          shard_ids, shard_unique_index = array_ops.unique(shard_ids)
-        with ops.device(self._impl.device(column)):
-          with Context.scope(sharding=False):
-            shard_embs = self._impl.lookup(
-              column, weights, shard_ids, sharded=True)
-        with ops.name_scope('shuffle_embeddings'):
-          if shard_unique_index is not None:
-            shard_embs = array_ops.gather(
-              shard_embs, shard_unique_index, name='restore_unique')
-          embs_shards = CommunicatorPool.get().call(
-            self._exchange_embeddings(dimension), [shard_embs, embs_sizes])
-          embeddings = array_ops.gather(
-            embs_shards, partition_index, name='restore_shuffle')
-      else:
-        with ops.device(self._impl.device(column)):
-          with Context.scope(sharding=False):
-            embeddings = self._impl.lookup(column, weights, ids)
-      return segment_reduce(
-        sparse_ids, embeddings,
-        weights=sparse_weights,
-        indices=unique_index,
-        dimension=dimension,
-        pad=pad,
-        segment_rank=segment_rank,
-        combiner=self._impl.combiner(column))
-
-
-class EmbeddingLookupCoalesced(object):  # pylint: disable=useless-object-inheritance
-  r'''Functor to lookup embeddings in group.
-  '''
-  def __init__(self):
-    self._impl = EmbeddingBackend.get()
-
-  def _exchange_group_ids(self, comm, inputs, inputs_deps):
-    r'''Communicator call to exchange group IDs.
-    '''
-    num_columns = len(inputs) // 2
-    if len(inputs) != 2 * num_columns:
-      raise ValueError('Number of IDs to exchange is wrong')
-    values = inputs[:num_columns]
-    sizes = inputs[num_columns:]
-    with ops.control_dependencies(inputs_deps):
-      return comm.alltoallv_n(values, sizes), None
-
-  def _exchange_group_embeddings(self, wire_dtype, common_shapes):
-    r'''Communicator call to exchange group embeddings.
-    '''
-    def _fn(comm, inputs, inputs_deps):
-      num_columns = len(inputs) // 2
-      if len(inputs) != 2 * num_columns:
-        raise ValueError('Number of IDs to exchange is wrong')
-      values = inputs[:num_columns]
-      sizes = inputs[num_columns:]
-      if inputs_deps is None:
-        embs, _ = comm.alltoallv_n(
-          values, sizes,
-          wire_dtype=wire_dtype,
-          common_shapes=common_shapes)
-        return embs, None
-      with ops.control_dependencies(inputs_deps):
-        embs, _ = comm.alltoallv_n(
-          values, sizes,
-          wire_dtype=wire_dtype,
-          common_shapes=common_shapes)
-
-      def grad_fn(grads, grads_deps, roots):
-        r'''Gradient function.
-        '''
-        del roots
-        with ops.control_dependencies(grads_deps):
-          d_inputs = gradients.gradients(embs, inputs, grad_ys=grads)
-        return d_inputs, None
-      return embs, grad_fn
-    return _fn
-
-  def __call__(
-      self, group_weights, group_inputs, columns, num_groups=1, name=None):
-    r'''Lookup embedding results.
-    '''
-    if name is None:
-      name = ops.get_default_graph().unique_name('embedding_lookup_coalesced')
-    num_shards = len(Context.get().devices)
-    num_columns = len(columns)
-    num_groups = min(num_groups, num_columns)
-    bucket_size = num_columns // num_groups
-    bucket_sizes = [bucket_size for _ in range(num_groups)]
-    for r in range(num_columns - num_groups * bucket_size):
-      bucket_sizes[r] += 1
-    bi = 0
-    bidx = []
-    for bucket_size in bucket_sizes:
-      bidx.append(bi)
-      bi += bucket_size
-    bidx.append(bi)
-
-    # Lookup embeddings in buckets.
-    group_embeddings = []
-    for bid, _ in enumerate(bucket_sizes):
-      with ops.name_scope(f'{name}/bucket{bid}'):
-        bucket_columns = columns[bidx[bid]: bidx[bid + 1]]
-        bucket_sparse_ids = []
-        bucket_sparse_weights = []
-        bucket_ids = []
-        bucket_unique_index = []
-        bucket_pad = []
-        with ops.name_scope('unique_ids'):
-          for idx, c in enumerate(bucket_columns):
-            inputs = group_inputs[bidx[bid]: bidx[bid + 1]][idx]
-            sparse_ids = inputs.id_tensor
-            bucket_sparse_ids.append(sparse_ids)
-            sparse_weights = inputs.weight_tensor
-            bucket_sparse_weights.append(sparse_weights)
-            if isinstance(sparse_ids, sparse_tensor.SparseTensor):
-              ids = sparse_ids.values
-            else:
-              ids = sparse_ids
-            with ops.device(self._impl.input_device(c)):
-              ids = array_ops.reshape(ids, [-1])
-              bucket_pad.append(self._impl.pad(c))
-              if self._impl.unique(c):
-                unique_index = None
-                if not bucket_pad[idx]:
-                  logging.info(
-                    'If unique is set to True, pad would be set to True')
-                  bucket_pad[idx] = True
-              else:
-                ids, unique_index = array_ops.unique(ids)
-              if self._impl.segment_rank(c) and not bucket_pad[idx]:
-                logging.info(
-                  'If segment rank is not 0, pad would be set to True')
-                bucket_pad[idx] = True
-              bucket_ids.append(ids)
-              bucket_unique_index.append(unique_index)
-
-        bucket_ids_shards, bucket_ids_sizes, bucket_partition_index = (
-          floormod_shuffle_n(bucket_ids, num_shards))
-        bucket_shard_ids, bucket_embs_sizes = CommunicatorPool.get().call(
-          self._exchange_group_ids,
-          bucket_ids_shards + bucket_ids_sizes,
-          trainable=False)
-
-        bucket_weights = group_weights[bidx[bid]: bidx[bid + 1]]
-        bucket_shard_embs = []
-        with ops.name_scope('unique_embs'):
-          for idx, c in enumerate(bucket_columns):
-            shard_ids = bucket_shard_ids[idx]
-            shard_ids, shard_unique_index = array_ops.unique(shard_ids)
-            with ops.device(self._impl.device(c)):
-              with Context.scope(sharding=False):
-                shard_embs = self._impl.lookup(
-                  c, bucket_weights[idx], shard_ids,
-                  sharded=self._impl.sharded(c))
-            if shard_unique_index is not None:
-              shard_embs = array_ops.gather(
-                shard_embs, shard_unique_index, name='restore_unique')
-            bucket_shard_embs.append(shard_embs)
-
-        common_shapes = [[self._impl.dimension(c)] for c in bucket_columns]
-        wire_dtypes = {self._impl.wire_dtype(c) for c in bucket_columns}
-        if len(wire_dtypes) > 1:
-          raise ValueError(
-            f'Multiple wire data type ({wire_dtypes}) not supported '
-            'in coalescing mode, use consistent `wire_dtype` or '
-            'set `num_groups` to `None`')
-        bucket_embs_shards = CommunicatorPool.get().call(
-          self._exchange_group_embeddings(list(wire_dtypes)[0], common_shapes),
-          bucket_shard_embs + bucket_embs_sizes)
-
-        with ops.name_scope('segment_reduce'):
-          for idx, c in enumerate(bucket_columns):
-            embeddings = array_ops.gather(
-              bucket_embs_shards[idx], bucket_partition_index[idx],
-              name='restore_shuffle')
-            embeddings = segment_reduce(
-              bucket_sparse_ids[idx], embeddings,
-              weights=bucket_sparse_weights[idx],
-              indices=bucket_unique_index[idx],
-              dimension=self._impl.dimension(c),
-              pad=bucket_pad[idx],
-              segment_rank=self._impl.segment_rank(c),
-              combiner=self._impl.combiner(c))
-            group_embeddings.append(embeddings)
-    return group_embeddings
diff --git a/hybridbackend/tensorflow/feature_column/math_lib.py b/hybridbackend/tensorflow/feature_column/math_lib.py
deleted file mode 100644
index dcd44b8a..00000000
--- a/hybridbackend/tensorflow/feature_column/math_lib.py
+++ /dev/null
@@ -1,613 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Arithmetic operators for dense and sparse tensors.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import math_ops
-
-try:
-  from tensorflow.python.ops import sort_ops
-except ImportError:
-  from tensorflow.contrib.framework.python.ops import sort_ops
-
-from hybridbackend.tensorflow.framework.context import Context
-
-
-class SegmentReduceException(Exception):
-  r'''handling error message from segment_reduction.
-  '''
-  def __init__(self, error_msg, pad, segment_rank):
-    super().__init__(self)
-    self._error_msg = error_msg
-    self._pad = pad
-    self._segment_rank = segment_rank
-
-  def __str__(self):
-    if self._pad:
-      self._error_msg += (
-        'pad of this column is set to True, either it is set via emb_pad or '
-        'the emb_unique is set to true or segment_rank is non zero. '
-        f'In any of these cases, make sure that dim[{self._segment_rank}] '
-        'of the dense_shape of ids must be specified')
-    return self._error_msg
-
-
-def _get_segment_ids_and_num(ids, segment_rank):
-  r'''Calculate segment ids and num_segments with a specifed segment_rank.
-  '''
-  if segment_rank < 0 or segment_rank > (ids.dense_shape.shape[0].value - 1):
-    raise ValueError(
-      f'invalid segment_rank {segment_rank}, must be within'
-      f' [0, {ids.dense_shape.shape[0].value})')
-  segment_ids = ids.indices[:, segment_rank]
-  stride = ids.dense_shape[segment_rank]
-  for i in reversed(range(segment_rank)):
-    segment_ids = math_ops.add_n(
-      [segment_ids, math_ops.multiply(ids.indices[:, i], stride)])
-    stride = stride * ids.dense_shape[i]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  if Context.get().options.emb_segment_sort:
-    segment_ids = sort_ops.sort(segment_ids)
-  return segment_ids, stride
-
-
-def _segment_sum(ids, data, weights=None, pad=False, segment_rank=0, name=None):
-  r'''Do segment sum with weights.
-
-  Example:
-    data: [1., 2., 3.]
-    weights: [2., 2., 2.]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 4]
-    pad: False
-    segment_rank: 0
-    output: [2, 10, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    weights: has the same length as data.
-    pad: whether or not padding the output.
-    segment_rank: segment reduce by values on this rank, default is 0.
-
-  Returns:
-    summation of data and weights by segments.
-  '''
-  if name is None:
-    name = 'segment_sum'
-
-  if weights is not None:
-    data = data * weights
-  if pad:
-    segment_ids, num_segments = _get_segment_ids_and_num(ids, segment_rank)
-    return math_ops.unsorted_segment_sum(
-      data, segment_ids, num_segments, name=name)
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  return math_ops.segment_sum(data, segment_ids, name=name)
-
-
-def _sparse_segment_sum(
-    ids, data, indices, pad=False, segment_rank=0, name=None):
-  r'''Segment sum with entries of data specifed by indices.
-
-  Example:
-    data: [1., 2., 3.]
-    indices: [0, 1]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 4]
-    pad: False
-    segment_rank: 0
-    output: [1, 2, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    indices: specify the entries in data to be summed up.
-    pad: whether or not padding the output.
-    segment_rank: segment reduce by values on this rank, default is 0
-
-  Returns:
-    summation of data by segments.
-  '''
-  if name is None:
-    name = 'sparse_segment_sum'
-
-  if pad:
-    segment_ids, num_segments = _get_segment_ids_and_num(ids, segment_rank)
-    return math_ops.sparse_segment_sum_with_num_segments(
-      data, indices, segment_ids, name=name, num_segments=num_segments)
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  return math_ops.sparse_segment_sum(
-    data, indices, segment_ids, name=name)
-
-
-def _segment_mean(
-    ids, data, weights=None, pad=False, segment_rank=0, name=None):
-  r'''Segment mean with weights.
-
-  Example:
-    data: [1., 2., 3.]
-    weights: [2., 2., 2.]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 4]
-    pad: True
-    segment_rank: 0
-    output: [1, 2.5, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    weights: has the same length as data.
-    pad: whether or not padding the output.
-    segment_rank: segment reduce by values on this rank, default is 0
-
-  Returns:
-    mean of segment summed weighted data by segment summed weights.
-  '''
-  if name is None:
-    name = 'segment_mean'
-
-  if weights is not None:
-    data = data * weights
-  if pad:
-    segment_ids, num_segments = _get_segment_ids_and_num(ids, segment_rank)
-    segmented = math_ops.unsorted_segment_sum(
-      data, segment_ids, num_segments)
-    if weights is None:
-      return segmented
-    weight_sum = math_ops.unsorted_segment_sum(
-      weights, segment_ids, num_segments)
-    return math_ops.div(segmented, weight_sum, name=name)
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  segmented = math_ops.segment_sum(data, segment_ids)
-  if weights is None:
-    return segmented
-  weight_sum = math_ops.segment_sum(weights, segment_ids)
-  return math_ops.div(segmented, weight_sum, name=name)
-
-
-def _sparse_segment_mean(
-    ids, data, indices, pad=False, segment_rank=0, name=None):
-  r'''Segment mean with entries of data specifed by indices.
-
-  Example:
-    data: [1., 2., 3.]
-    indices: [0, 1]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 4]
-    pad: True
-    segment_rank: 0
-    output: [0.25, 0.5, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    indices: specify the entries in data to be summed up.
-    pad: whether or not padding the output.
-    segment_rank: segment reduce by values on this rank, default is 0
-
-  Returns:
-    mean of segment summed data by segments.
-  '''
-  if name is None:
-    name = 'sparse_segment_mean'
-
-  if pad:
-    segment_ids, num_segments = _get_segment_ids_and_num(ids, segment_rank)
-    return math_ops.sparse_segment_mean_with_num_segments(
-      data, indices, segment_ids, name=name, num_segments=num_segments)
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  return math_ops.sparse_segment_mean(
-    data, indices, segment_ids, name=name)
-
-
-def _segment_sqrtn(
-    ids, data, weights=None, pad=False, segment_rank=0, name=None):
-  r'''Segment sqrtn with weights.
-
-  Example:
-    data: [1., 2., 3.]
-    weights: [2., 2., 2.]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 4]
-    pad: True
-    segment_rank: 0
-    output: [1, 3.57, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    weights: has the same length as data.
-    pad: whether or not padding the output.
-    segment_rank: segment reduce by values on this rank, default is 0
-
-  Returns:
-    sqrtn of segment summed weighted data by segment summed weights.
-  '''
-  if name is None:
-    name = 'segment_sqrtn'
-
-  if weights is not None:
-    data = data * weights
-  if pad:
-    segment_ids, num_segments = _get_segment_ids_and_num(ids, segment_rank)
-    segmented = math_ops.unsorted_segment_sum(
-      data, segment_ids, num_segments)
-    if weights is None:
-      return segmented
-    weights_squared = math_ops.pow(weights, 2)
-    weight_sum = math_ops.unsorted_segment_sum(
-      weights_squared, segment_ids, num_segments)
-    weight_sum_sqrt = math_ops.sqrt(weight_sum)
-    return math_ops.div(segmented, weight_sum_sqrt, name=name)
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  segmented = math_ops.segment_sum(data, segment_ids)
-  if weights is None:
-    return segmented
-  weights_squared = math_ops.pow(weights, 2)
-  weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
-  weight_sum_sqrt = math_ops.sqrt(weight_sum)
-  return math_ops.div(segmented, weight_sum_sqrt, name=name)
-
-
-def _sparse_segment_sqrtn(
-    ids, data, indices, pad=False, segment_rank=0, name=None):
-  r'''Segment sqrtn with entries of data specifed by indices.
-
-  Example:
-    data: [1., 2., 3.]
-    indices: [0, 1]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 4]
-    pad: True
-    segment_rank: 0
-    output: [0.5, 1, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    indices: specify the entries in data to be summed up.
-    pad: whether or not padding the output.
-    segment_rank: segment reduce by values on this rank, default is 0
-
-  Returns:
-    sqrtn of segment summed data by segments.
-  '''
-  if name is None:
-    name = 'sparse_segment_sqrtn'
-
-  if pad:
-    segment_ids, num_segments = _get_segment_ids_and_num(ids, segment_rank)
-    return math_ops.sparse_segment_sqrt_n_with_num_segments(
-      data, indices, segment_ids, name=name, num_segments=num_segments)
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  return math_ops.sparse_segment_sqrt_n(
-    data, indices, segment_ids, name=name)
-
-
-def _segment_tile(ids, data, weights=None, pad=False, name=None):
-  r'''Segment tiling with weights.
-
-  Example:
-    data: [1., 2., 3.]
-    weights: [2., 2., 2.]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 2]
-    pad: True
-    num_tiles = 2
-    tile_ids = [1, 2, 3]
-    tiled_segment_ids = [1, 4, 5]
-    segments = 8
-    weighted_data = [2, 4, 6]
-    output = [0, 2, 0, 0, 4, 6, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    weights: has the same length as data.
-    pad: whether or not padding the output.
-
-  Returns:
-    tiling of segment summed weighted data.
-  '''
-  if name is None:
-    name = 'segment_tile'
-
-  dim = data.shape[1]
-  if weights is not None:
-    data = data * weights
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  tile_ids = ids.indices[:, 1]
-  if tile_ids.dtype != dtypes.int32:
-    tile_ids = math_ops.cast(tile_ids, dtypes.int32)
-  num_segments = ids.dense_shape[0]
-  if num_segments.dtype != dtypes.int32:
-    num_segments = math_ops.cast(num_segments, dtypes.int32)
-  num_tiles = ids.dense_shape[1]
-  if num_tiles.dtype != dtypes.int32:
-    num_tiles = math_ops.cast(num_tiles, dtypes.int32)
-  tiled_segment_ids = segment_ids * num_tiles + tile_ids
-  if pad:
-    segmented = math_ops.unsorted_segment_sum(
-      data, tiled_segment_ids, num_segments * num_tiles, name=name)
-  else:
-    segmented = math_ops.segment_sum(
-      data, tiled_segment_ids, name=name)
-  return array_ops.reshape(segmented, [-1, num_tiles * dim])
-
-
-def _sparse_segment_tile(ids, data, indices, pad=False, name=None):
-  r'''Segment tiling with specified indices.
-
-  Example:
-    data: [1., 2., 3.]
-    indices: [0, 1]
-    ids.indices: [0, 1], [1, 2], [1, 3]
-    ids.dense_shape: [4, 2]
-    pad: False
-    num_tiles = 2
-    tile_ids = [1, 2, 3]
-    tiled_segment_ids = [1, 4, 5]
-    segments = 8
-    output = [0, 1, 0, 0, 4, 0, 0, 0]
-
-  Args:
-    ids: sparse_tensor.SparseTensor, has the same length as data.
-    data: embedding values.
-    indices: specify the entries in data to be summed up.
-    pad: whether or not padding the output.
-
-  Returns:
-    tiling of segment summed data by segments.
-  '''
-  if name is None:
-    name = 'sparse_segment_tile'
-
-  segment_ids = ids.indices[:, 0]
-  if segment_ids.dtype != dtypes.int32:
-    segment_ids = math_ops.cast(segment_ids, dtypes.int32)
-  tile_ids = ids.indices[:, 1]
-  num_tiles = ids.dense_shape[1]
-  tiled_segment_ids = segment_ids * num_tiles + tile_ids
-  if pad:
-    num_segments = ids.dense_shape[0]
-    return math_ops.sparse_segment_sum_with_num_segments(
-      data, indices, tiled_segment_ids,
-      name=name,
-      num_segments=num_segments * num_tiles)
-  return math_ops.sparse_segment_sum(
-    data, indices, tiled_segment_ids, name=name)
-
-
-def _segment_reduce(
-    ids, data,
-    weights=None,
-    pad=False,
-    segment_rank=0,
-    combiner=None,
-    name=None):
-  r'''Computes the sum along segments of a tensor.
-
-  Args:
-    ids: A `SparseTensor` as original input.
-    data: A `Tensor` with data that will be assembled in the output.
-    weights: Weights for data.
-    pad: If True, pad the output.
-    segment_rank: segment reduce by values on this rank, default is 0
-    combiner: A string specifying the reduction op.
-    name: Optional name of the operation.
-  '''
-  combiner = combiner or 'sum'
-  with ops.name_scope(combiner):
-    if combiner == 'sum':
-      return _segment_sum(
-        ids, data, weights, pad=pad, segment_rank=segment_rank, name=name)
-    if combiner == 'mean':
-      return _segment_mean(
-        ids, data, weights, pad=pad, segment_rank=segment_rank, name=name)
-    if combiner == 'sqrtn':
-      return _segment_sqrtn(
-        ids, data, weights, pad=pad, segment_rank=segment_rank, name=name)
-    if combiner == 'tile':
-      return _segment_tile(ids, data, weights, pad=pad, name=name)
-    if callable(combiner):
-      if name is None:
-        name = 'segment_combine'
-      return combiner(
-        ids, data,
-        weights=weights,
-        indices=None,
-        pad=pad,
-        name=name)
-  raise ValueError(f'Unrecognized combiner: {combiner}')
-
-
-def _sparse_segment_reduce(
-    ids, data, indices,
-    pad=False,
-    segment_rank=0,
-    combiner=None,
-    name=None):
-  r'''Computes the sum along sparse segments of a tensor.
-
-  Args:
-    ids: A `SparseTensor` as original input.
-    data: A `Tensor` with data that will be assembled in the output.
-    indices: A 1-D `Tensor` with indices into `data`. Has same rank as
-      `segment_ids`.
-    pad: If True, pad the output.
-    segment_rank: segment reduce by values on this rank, default is 0
-    combiner: A string specifying the reduction op.
-    name: Optional name of the operation.
-  '''
-  combiner = combiner or 'sum'
-  ids_shape = ids.dense_shape
-  ids_ = sparse_tensor.SparseTensor(ids.indices, ids.values, ids_shape)
-  with ops.name_scope(combiner):
-    if combiner == 'sum':
-      return _sparse_segment_sum(
-        ids_, data, indices,
-        pad=pad,
-        segment_rank=segment_rank,
-        name=name)
-    if combiner == 'mean':
-      return _sparse_segment_mean(
-        ids_, data, indices,
-        pad=pad,
-        segment_rank=segment_rank,
-        name=name)
-    if combiner == 'sqrtn':
-      return _sparse_segment_sqrtn(
-        ids_, data, indices,
-        pad=pad,
-        segment_rank=segment_rank,
-        name=name)
-    if combiner == 'tile':
-      return _sparse_segment_tile(
-        ids_, data, indices,
-        pad=pad,
-        name=name)
-    if callable(combiner):
-      return combiner(
-        ids_, data,
-        weights=None,
-        indices=indices,
-        pad=pad,
-        name=name)
-    raise ValueError(f'Unrecognized combiner: {combiner}')
-
-
-def segment_reduce(
-    ids, data,
-    weights=None,
-    indices=None,
-    dimension=None,
-    pad=False,
-    segment_rank=0,
-    combiner=None,
-    name=None):
-  r'''Computes the sum along segments of a tensor.
-
-  Args:
-    ids: A `SparseTensor` as original input.
-    data: A `Tensor` with data that will be assembled in the output.
-    weights: Weights with same shape of data.
-    indices: A 1-D `Tensor` with indices into `data`.
-    dimension: Dimension of embedding weights.
-    pad: If True, pad the output.
-    segment_rank: segment reduce by values on this rank, default is 0.
-    combiner: A string specifying the reduction op.
-    name: Optional name of the operation.
-  '''
-  if data.dtype in (dtypes.float16, dtypes.bfloat16):
-    data = math_ops.cast(data, dtypes.float32)
-
-  if not isinstance(ids, sparse_tensor.SparseTensor):
-    # Reshape to reverse the flattening of ids.
-    with ops.name_scope('update_shape'):
-      dims = tensor_shape.TensorShape([dimension])
-      embeddings_shape = ids.get_shape().concatenate(dims)
-      ids_shape = array_ops.shape(ids, name='ids_shape')
-      if dims.is_fully_defined():
-        actual_embedding_shape = array_ops.concat(
-          [ids_shape, dims], 0, name='actual_embedding_shape')
-      else:
-        actual_embedding_shape = array_ops.concat(
-          [ids_shape, array_ops.shape(data)[1:]], 0,
-          name='actual_embedding_shape')
-    # Recover unique
-    if indices is not None:
-      data = array_ops.gather(data, indices, name='restore_unique')
-    embeddings = array_ops.reshape(
-      data, actual_embedding_shape, name=name)
-    embeddings.set_shape(embeddings_shape)
-    return embeddings
-
-  try:
-    if weights is not None:
-      if indices is not None:
-        data = array_ops.gather(data, indices, name='restore_unique')
-      weights = weights.values
-
-      if weights.dtype != data.dtype:
-        weights = math_ops.cast(weights, data.dtype)
-
-      # Reshape weights to allow broadcast
-      ones = array_ops.fill(
-        array_ops.expand_dims(array_ops.rank(data) - 1, 0), 1)
-      bcast_weights_shape = array_ops.concat(
-        [array_ops.shape(weights), ones], 0)
-      orig_weights_shape = weights.get_shape()
-      weights = array_ops.reshape(weights, bcast_weights_shape)
-
-      # Set the weight shape, since after reshaping to
-      # bcast_weights_shape, the shape becomes None.
-      if data.get_shape().ndims is not None:
-        weights.set_shape(
-          orig_weights_shape.concatenate(
-            [1 for _ in range(data.get_shape().ndims - 1)]))
-
-      embeddings = _segment_reduce(
-        ids, data,
-        weights=weights,
-        pad=pad,
-        segment_rank=segment_rank,
-        combiner=combiner,
-        name=name)
-      return embeddings
-    if indices is not None:
-      embeddings = _sparse_segment_reduce(
-        ids, data, indices,
-        pad=pad,
-        segment_rank=segment_rank,
-        combiner=combiner,
-        name=name)
-      return embeddings
-    embeddings = _segment_reduce(
-      ids, data,
-      weights=None,
-      pad=pad,
-      segment_rank=segment_rank,
-      combiner=combiner,
-      name=name)
-  except ValueError as e:
-    raise SegmentReduceException(
-      str(e), pad, segment_rank) from e
-  return embeddings
diff --git a/hybridbackend/tensorflow/graph/common/helper.cc b/hybridbackend/tensorflow/graph/common/helper.cc
new file mode 100644
index 00000000..3f27a97b
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/helper.cc
@@ -0,0 +1,110 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include "hybridbackend/tensorflow/graph/common/helper.h"
+#include <absl/strings/str_join.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+namespace {
+struct GraphNodeFormatter {
+  void operator()(std::string* out, Node* n) const { out->append(n->name()); }
+};
+}  // namespace
+
+#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) < 1014L
+
+template <typename T>
+void DFSFromHelper(const Graph& g, gtl::ArraySlice<T> start,
+                   const std::function<void(T)>& enter,
+                   const std::function<void(T)>& leave,
+                   const NodeComparator& stable_comparator,
+                   const EdgeFilter& edge_filter) {
+  // Stack of work to do.
+  struct Work {
+    T node;
+    bool leave;  // Are we entering or leaving n?
+  };
+  std::vector<Work> stack(start.size());
+  for (int i = 0; i < start.size(); ++i) {
+    stack[i] = Work{start[i], false};
+  }
+
+  std::vector<bool> visited(g.num_node_ids(), false);
+  while (!stack.empty()) {
+    Work w = stack.back();
+    stack.pop_back();
+
+    T n = w.node;
+    if (w.leave) {
+      leave(n);
+      continue;
+    }
+
+    if (visited[n->id()]) continue;
+    visited[n->id()] = true;
+    if (enter) enter(n);
+
+    // Arrange to call leave(n) when all done with descendants.
+    if (leave) stack.push_back(Work{n, true});
+
+    auto add_work = [&visited, &stack](Node* out) {
+      if (!visited[out->id()]) {
+        // Note; we must not mark as visited until we actually process it.
+        stack.push_back(Work{out, false});
+      }
+    };
+
+    if (stable_comparator) {
+      std::vector<Node*> nodes_sorted;
+      for (const Edge* out_edge : n->out_edges()) {
+        if (!edge_filter || edge_filter(*out_edge)) {
+          nodes_sorted.emplace_back(out_edge->dst());
+        }
+      }
+      std::sort(nodes_sorted.begin(), nodes_sorted.end(), stable_comparator);
+      for (Node* out : nodes_sorted) {
+        add_work(out);
+      }
+    } else {
+      for (const Edge* out_edge : n->out_edges()) {
+        if (!edge_filter || edge_filter(*out_edge)) {
+          add_work(out_edge->dst());
+        }
+      }
+    }
+  }
+}
+
+void DFSFrom(const Graph& g, gtl::ArraySlice<Node*> start,
+             const std::function<void(Node*)>& enter,
+             const std::function<void(Node*)>& leave,
+             const NodeComparator& stable_comparator,
+             const EdgeFilter& edge_filter) {
+  DFSFromHelper(g, start, enter, leave, stable_comparator, edge_filter);
+}
+
+#endif
+
+string NodeJoin(const std::vector<Node*>& nodes, const string& delim) {
+  return absl::StrJoin(nodes, delim, GraphNodeFormatter());
+}
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/common/helper.h b/hybridbackend/tensorflow/graph/common/helper.h
new file mode 100644
index 00000000..dae6c7ee
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/helper.h
@@ -0,0 +1,52 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_HELPER_H_
+#define HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_HELPER_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <string>
+#include <vector>
+
+#include <tensorflow/core/graph/algorithm.h>
+#include <tensorflow/core/graph/graph.h>
+#include <tensorflow/core/public/version.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+#if (TF_MAJOR_VERSION * 1000L + TF_MINOR_VERSION) < 1014L
+template <typename T>
+void DFSFromHelper(const Graph& g, gtl::ArraySlice<T> start,
+                   const std::function<void(T)>& enter,
+                   const std::function<void(T)>& leave,
+                   const NodeComparator& stable_comparator,
+                   const EdgeFilter& edge_filter);
+
+void DFSFrom(const Graph& g, gtl::ArraySlice<Node*> start,
+             const std::function<void(Node*)>& enter,
+             const std::function<void(Node*)>& leave,
+             const NodeComparator& stable_comparator = {},
+             const EdgeFilter& edge_filter = {});
+#endif
+
+string NodeJoin(const std::vector<Node*>& nodes, const string& delim);
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
+#endif  // HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_HELPER_H_
diff --git a/hybridbackend/tensorflow/graph/common/linearization.cc b/hybridbackend/tensorflow/graph/common/linearization.cc
new file mode 100644
index 00000000..39a5e18e
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/linearization.cc
@@ -0,0 +1,87 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <algorithm>
+#include <deque>
+#include <set>
+#include <unordered_map>
+#include <vector>
+
+#include <absl/strings/str_cat.h>
+#include <absl/strings/str_join.h>
+#include <absl/strings/str_split.h>
+
+#include <tensorflow/core/graph/node_builder.h>
+#include <tensorflow/core/public/version.h>
+#include <tensorflow/core/util/device_name_utils.h>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/helper.h"
+#include "hybridbackend/tensorflow/graph/common/linearization.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+LinearizeOutputs::LinearizeOutputs(const string& op_type,
+                                   const int32& op_output)
+    : op_type_(op_type), op_output_(op_output) {}
+
+Status LinearizeOutputs::In(Graph* graph) {
+  std::unordered_map<Node*, int> candidates;
+  std::vector<bool> dependencies;
+
+  std::vector<Node*> sorted;
+  GetReversePostOrder(*graph, &sorted, NodeComparatorName{});
+
+  for (Node* node : graph->op_nodes()) {
+    if (!node->IsOp()) {
+      continue;
+    }
+
+    if (node->type_string() != op_type_) {
+      continue;
+    }
+
+    std::vector<Node*> linear_ops;
+    for (Node* n : sorted) {
+      for (const auto& edge : n->in_edges()) {
+        if (edge && !edge->IsControlEdge() && edge->src() == node &&
+            edge->src_output() == op_output_) {
+          linear_ops.push_back(edge->dst());
+          break;
+        }
+      }
+    }
+
+    if (linear_ops.size() < 2) {
+      continue;
+    }
+
+    for (size_t idx = 1; idx < linear_ops.size(); ++idx) {
+      graph->AddControlEdge(linear_ops[idx - 1], linear_ops[idx]);
+    }
+
+    VLOG(1) << "Linearized " << linear_ops.size() << " outputs of "
+            << node->name() << " in graph " << static_cast<void*>(graph);
+  }
+
+  return Status::OK();
+}
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/common/linearization.h b/hybridbackend/tensorflow/graph/common/linearization.h
new file mode 100644
index 00000000..d5cda1b5
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/linearization.h
@@ -0,0 +1,46 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_LINEARIZATION_H_
+#define HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_LINEARIZATION_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include <tensorflow/core/graph/graph.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+class LinearizeOutputs {
+ public:
+  LinearizeOutputs(const string& op_type, const int32& op_output);
+  Status In(Graph* graph);
+
+ private:
+  string op_type_;
+  int32 op_output_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(LinearizeOutputs);
+};
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
+#endif  // HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_LINEARIZATION_H_
diff --git a/hybridbackend/tensorflow/graph/common/packing.cc b/hybridbackend/tensorflow/graph/common/packing.cc
new file mode 100644
index 00000000..54b429bd
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/packing.cc
@@ -0,0 +1,580 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <algorithm>
+#include <deque>
+#include <sstream>
+#include <unordered_map>
+#include <vector>
+
+#include <absl/strings/str_cat.h>
+#include <absl/strings/str_join.h>
+#include <absl/strings/str_split.h>
+
+#include <tensorflow/core/graph/node_builder.h>
+#include <tensorflow/core/public/version.h>
+#include <tensorflow/core/util/device_name_utils.h>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/helper.h"
+#include "hybridbackend/tensorflow/graph/common/packing.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+namespace {
+inline bool PackingDisabled() {
+  static const bool kPackingDisabled =
+      ::hybridbackend::EnvVarGetBool("HB_OP_PACKING_DISABLED", false);
+  return kPackingDisabled;
+}
+
+inline string ClustersDebugString(
+    const std::unordered_map<string, std::vector<Node*>>& clusters) {
+  std::ostringstream debug_sstr;
+  size_t cluster_size = clusters.size();
+  size_t idx = 0;
+  for (auto& c : clusters) {
+    debug_sstr << c.first << " (" << c.second.size() << ")";
+    if (idx < cluster_size - 1) {
+      debug_sstr << ", ";
+    }
+    ++idx;
+  }
+  return debug_sstr.str();
+}
+}  // namespace
+
+Pack::Pack(const string& op_type, const string& optimized_op_type)
+    : op_type_(op_type),
+      optimized_op_type_(optimized_op_type),
+      device_(DEVICE_GPU) {
+  static const int kDefaultPackingNBuckets =
+      ::hybridbackend::EnvVarGetInt("HB_OP_PACKING_NBUCKETS", 1);
+  num_buckets_ = kDefaultPackingNBuckets;
+}
+
+Pack& Pack::WithDevice(const string& device) {
+  device_ = device;
+  return *this;
+}
+
+Pack& Pack::WithTypeAttr(const string& attr_name,
+                         const std::vector<DataType>& constraints) {
+  type_attrs_[attr_name] = constraints;
+  return *this;
+}
+
+Pack& Pack::WithShapeAttr(const string& attr_name) {
+  shape_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Pack& Pack::WithIntAttr(const string& attr_name) {
+  int_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Pack& Pack::WithStrAttr(const string& attr_name) {
+  str_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Pack& Pack::WithAggregatedShapeAttr(const string& attr_name) {
+  aggregated_shape_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Pack& Pack::WithAggregatedIntAttr(const string& attr_name) {
+  aggregated_int_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Pack& Pack::WithAggregatedStrAttr(const string& attr_name) {
+  aggregated_str_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Pack& Pack::WithHandle(const int32 input) {
+  handles_.push_back(input);
+  return *this;
+}
+
+Pack& Pack::WithBuckets(const int32 num_buckets) {
+  num_buckets_ = num_buckets;
+  return *this;
+}
+
+Status Pack::In(Graph* graph) { return In(graph, nullptr); }
+
+Status Pack::In(Graph* graph, int64* poccurrence_count) {
+  int64 occurrence_count = 0;
+  std::unordered_map<Node*, int> candidates;
+  std::vector<bool> dependencies;
+
+  std::vector<Node*> sorted;
+  GetReversePostOrder(*graph, &sorted);
+
+  std::vector<Node*> roots;
+  int seq = 0;
+  for (Node* node : sorted) {
+    if (!node->IsOp()) {
+      continue;
+    }
+
+    if (node->type_string() != op_type_) {
+      continue;
+    }
+
+    const int32 num_inputs = node->num_inputs();
+    if (num_inputs < 1) {
+      VLOG(2) << "Skipped packing of op " << node->name() << " since no inputs";
+      continue;
+    }
+
+    DeviceNameUtils::ParsedName requested_device;
+    if (!DeviceNameUtils::ParseFullName(node->requested_device(),
+                                        &requested_device)) {
+      VLOG(2) << "Skipped packing of op " << node->name()
+              << " since no device set";
+      continue;
+    }
+
+    if (requested_device.type != device_ || requested_device.job == "ps") {
+      VLOG(2) << "Skipped packing of op " << node->name()
+              << " since device differs: /" << device_ << ":0 vs. "
+              << node->requested_device();
+      continue;
+    }
+
+    bool types_are_incompatible = false;
+    for (auto& attr_pair : type_attrs_) {
+      if (!HasNodeAttr(node->def(), attr_pair.first)) {
+        VLOG(2) << "Skipped packing of op " << node->name() << " since "
+                << op_type_ << " nodes has no attribute " << attr_pair.first;
+        types_are_incompatible = true;
+        break;
+      }
+      DataType attrval;
+      TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attr_pair.first, &attrval));
+      if (attr_pair.second.size() == 0) {
+        continue;
+      }
+      types_are_incompatible = true;
+      for (DataType dt : attr_pair.second) {
+        if (dt == attrval) {
+          types_are_incompatible = false;
+          break;
+        }
+      }
+      if (TF_PREDICT_FALSE(types_are_incompatible)) {
+        VLOG(2) << "Skipped packing of op " << node->name() << " since "
+                << op_type_ << " nodes with attribute " << attr_pair.first
+                << "=" << DataTypeString(attrval) << " is not supported";
+      }
+    }
+
+    if (TF_PREDICT_FALSE(types_are_incompatible)) {
+      continue;
+    }
+
+    roots.push_back(node);
+    candidates.emplace(node, seq++);
+  }
+  const size_t size = candidates.size();
+  if (size == 0) {
+    return Status::OK();
+  }
+
+  dependencies.resize(size * size, false);
+  for (size_t i = 0; i < size; ++i) {
+    dependencies[i * size + i] = true;
+  }
+
+  for (size_t i = 0; i < size; ++i) {
+    Node* root = roots[i];
+    auto root_it = candidates.find(root);
+    DFSFrom(*graph, {root},
+            [&](Node* n) {
+              auto it = candidates.find(n);
+              if (it != candidates.end()) {
+                dependencies[root_it->second * size + it->second] = true;
+              }
+            },
+            nullptr);
+  }
+
+  if (TF_PREDICT_FALSE(PackingDisabled())) {
+    return Status::OK();
+  }
+
+  std::vector<int> colors(size, 0);
+  for (int color = 0; color < size; ++color) {
+    int cursor = 0;
+    bool color_is_set = false;
+    for (; cursor < size; ++cursor) {
+      if (colors[cursor] != color) {
+        continue;
+      }
+      for (size_t i = cursor + 1; i < size; ++i) {
+        if (dependencies[cursor * size + i]) {
+          colors[i] = color + 1;
+        }
+      }
+      color_is_set = true;
+    }
+    if (!color_is_set) {
+      break;
+    }
+  }
+
+  std::unordered_map<string, std::vector<Node*>> scoped_clusters;
+  for (auto candidate : candidates) {
+    Node* node = candidate.first;
+
+    std::vector<string> attrs;
+    for (auto& attr_pair : type_attrs_) {
+      if (HasNodeAttr(node->def(), attr_pair.first)) {
+        DataType attrval;
+        TF_RETURN_IF_ERROR(
+            GetNodeAttr(node->attrs(), attr_pair.first, &attrval));
+        attrs.push_back(DataTypeString(attrval));
+      } else {
+        attrs.push_back("_");
+      }
+    }
+    for (auto& attr : shape_attrs_) {
+      if (HasNodeAttr(node->def(), attr)) {
+        TensorShape attrval;
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attr, &attrval));
+        attrs.push_back(attrval.DebugString());
+      } else {
+        attrs.push_back("_");
+      }
+    }
+    for (auto& attr : int_attrs_) {
+      if (HasNodeAttr(node->def(), attr)) {
+        int attrval;
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attr, &attrval));
+        attrs.push_back(std::to_string(attrval));
+      } else {
+        attrs.push_back("_");
+      }
+    }
+    for (auto& attr : str_attrs_) {
+      if (HasNodeAttr(node->def(), attr)) {
+        string attrval;
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attr, &attrval));
+        attrs.push_back(attrval);
+      } else {
+        attrs.push_back("_");
+      }
+    }
+
+    std::vector<string> handles;
+    for (auto& idx : handles_) {
+      const Edge* input_edge;
+      TF_RETURN_IF_ERROR(node->input_edge(idx, &input_edge));
+      handles.push_back(input_edge->src()->name());
+    }
+
+    auto name = node->name();
+    auto prefix = name.substr(0, name.find_first_of('/'));
+
+    int seq = candidates[node];
+    string cluster_key =
+        absl::StrCat(colors[seq], ":", absl::StrJoin(attrs, "|"), ":",
+                     absl::StrJoin(handles, "|"), ":", prefix);
+    auto it = scoped_clusters.find(cluster_key);
+    if (it == scoped_clusters.end()) {
+      scoped_clusters.emplace(cluster_key, std::vector<Node*>{node});
+    } else {
+      it->second.push_back(node);
+    }
+  }
+
+  std::unordered_map<string, std::vector<Node*>> clusters;
+  for (auto& c : scoped_clusters) {
+    int cluster_size = c.second.size();
+    if (cluster_size < 2) {
+      continue;
+    }
+    if (cluster_size < num_buckets_ * 2) {
+      clusters.emplace(c.first, c.second);
+      continue;
+    }
+    std::vector<std::vector<Node*>> buckets(num_buckets_);
+    for (size_t i = 0; i < cluster_size; ++i) {
+      buckets[i % num_buckets_].push_back(c.second[i]);
+    }
+    for (size_t b = 0; b < num_buckets_; ++b) {
+      clusters.emplace(absl::StrCat(c.first, ":", b), buckets[b]);
+    }
+  }
+
+  VLOG(3) << "Try pack " << size << " " << op_type_ << " ops into "
+          << clusters.size() << " clusters: " << ClustersDebugString(clusters);
+
+  for (auto& c : clusters) {
+    if (c.second.size() <= 1) {
+      continue;
+    }
+
+    Node* node0 = c.second[0];
+    std::unordered_map<string, AttrValue> node0attrs;
+    for (const auto& iter : node0->attrs()) {
+      node0attrs[iter.first] = iter.second;
+    }
+    bool all_attrs_are_equal = true;
+    for (size_t i = 1; i < c.second.size(); ++i) {
+      if (node0->num_inputs() != c.second[i]->num_inputs()) {
+        VLOG(2) << "Skipped packing of " << op_type_ << " ops (" << c.first
+                << ") with " << c.second.size()
+                << " nodes since num_inputs differs: " << node0->num_inputs()
+                << " vs. " << c.second[i]->num_inputs();
+        all_attrs_are_equal = false;
+        break;
+      }
+      if (node0->num_outputs() != c.second[i]->num_outputs()) {
+        VLOG(2) << "Skipped packing of " << op_type_ << " ops (" << c.first
+                << ") with " << c.second.size()
+                << " nodes since num_outputs differs: " << node0->num_outputs()
+                << " vs. " << c.second[i]->num_outputs();
+        all_attrs_are_equal = false;
+        break;
+      }
+      if (node0->requested_device() != c.second[i]->requested_device()) {
+        all_attrs_are_equal = false;
+        VLOG(2) << "Skipped packing of " << op_type_ << " ops (" << c.first
+                << ") with " << c.second.size()
+                << " nodes since requested device differs: "
+                << node0->requested_device() << " vs. "
+                << c.second[i]->requested_device();
+        break;
+      }
+      for (const auto& iter : c.second[i]->attrs()) {
+        if (iter.first == "_class") {
+          continue;
+        }
+        if (std::find(aggregated_shape_attrs_.begin(),
+                      aggregated_shape_attrs_.end(),
+                      iter.first) != aggregated_shape_attrs_.end()) {
+          continue;
+        }
+        if (std::find(aggregated_int_attrs_.begin(),
+                      aggregated_int_attrs_.end(),
+                      iter.first) != aggregated_int_attrs_.end()) {
+          continue;
+        }
+        if (std::find(aggregated_str_attrs_.begin(),
+                      aggregated_str_attrs_.end(),
+                      iter.first) != aggregated_str_attrs_.end()) {
+          continue;
+        }
+        auto iter0 = node0attrs.find(iter.first);
+        if (iter0 == node0attrs.end()) {
+          all_attrs_are_equal = false;
+          VLOG(2) << "Skipped packing of " << op_type_ << " ops (" << c.first
+                  << ") with " << c.second.size()
+                  << " nodes since attribute not found: " << iter.first;
+          break;
+        }
+        if (!AreAttrValuesEqual(iter.second, iter0->second)) {
+          all_attrs_are_equal = false;
+          VLOG(2) << "Skipped packing of " << op_type_ << " ops (" << c.first
+                  << ") with " << c.second.size() << " nodes since attribute "
+                  << iter.first << " differs";
+          break;
+        }
+      }
+      if (!all_attrs_are_equal) {
+        break;
+      }
+    }
+    if (!all_attrs_are_equal) {
+      continue;
+    }
+
+    std::sort(c.second.begin(), c.second.end(), NodeComparatorName{});
+
+    std::vector<std::vector<string>> fragments;
+    for (size_t n = 0; n < c.second.size(); ++n) {
+      std::vector<string> splits = absl::StrSplit(c.second[n]->name(), '/');
+      for (size_t i = 0; i < splits.size(); ++i) {
+        if (i < fragments.size()) {
+          fragments[i].push_back(splits[i]);
+        } else {
+          std::vector<string> new_splits{splits[i]};
+          fragments.emplace_back(std::move(new_splits));
+        }
+      }
+    }
+
+    std::vector<string> new_node_fragments;
+    for (auto& s : fragments) {
+      std::sort(s.begin(), s.end());
+      auto last = std::unique(s.begin(), s.end());
+      s.erase(last, s.end());
+      if (s.size() < 2) {
+        new_node_fragments.push_back(s[0]);
+        continue;
+      }
+      new_node_fragments.push_back(
+          absl::StrCat(s[0], "__", s.size(), "_packed"));
+    }
+    const string new_node_name =
+        graph->NewName(absl::StrJoin(new_node_fragments, "/"));
+    NodeBuilder packed_op_builder =
+        NodeBuilder(new_node_name, optimized_op_type_);
+    packed_op_builder.Device(node0->requested_device());
+    packed_op_builder.AssignedDevice(node0->assigned_device_name());
+    for (const auto& iter : node0->attrs()) {
+      if (iter.first == "N") {
+        continue;
+      }
+      if (std::find(aggregated_shape_attrs_.begin(),
+                    aggregated_shape_attrs_.end(),
+                    iter.first) != aggregated_shape_attrs_.end()) {
+        continue;
+      }
+      if (std::find(aggregated_int_attrs_.begin(), aggregated_int_attrs_.end(),
+                    iter.first) != aggregated_int_attrs_.end()) {
+        continue;
+      }
+      if (std::find(aggregated_str_attrs_.begin(), aggregated_str_attrs_.end(),
+                    iter.first) != aggregated_str_attrs_.end()) {
+        continue;
+      }
+      packed_op_builder.Attr(iter.first, iter.second);
+    }
+    packed_op_builder.Attr("N", static_cast<int>(c.second.size()));
+    for (const auto& attr_name : aggregated_shape_attrs_) {
+      std::vector<TensorShape> attr_values;
+      for (Node* n : c.second) {
+        TensorShape attrval;
+        TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &attrval));
+        attr_values.push_back(attrval);
+      }
+      packed_op_builder.Attr(attr_name, attr_values);
+    }
+    for (const auto& attr_name : aggregated_int_attrs_) {
+      std::vector<int32> attr_values;
+      for (Node* n : c.second) {
+        int32 attrval;
+        TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &attrval));
+        attr_values.push_back(attrval);
+      }
+      packed_op_builder.Attr(attr_name, attr_values);
+    }
+    for (const auto& attr_name : aggregated_str_attrs_) {
+      std::vector<string> attr_values;
+      for (Node* n : c.second) {
+        string attrval;
+        TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &attrval));
+        attr_values.push_back(attrval);
+      }
+      packed_op_builder.Attr(attr_name, attr_values);
+    }
+
+    std::vector<Node*> control_inputs;
+    for (Node* node : c.second) {
+      for (const auto& edge : node->in_edges()) {
+        if (!edge) {
+          continue;
+        }
+        if (edge->IsControlEdge()) {
+          control_inputs.push_back(edge->src());
+        }
+      }
+    }
+    packed_op_builder.ControlInputs(control_inputs);
+
+    std::vector<bool> handles_mask(node0->num_inputs(), false);
+    for (auto& idx : handles_) {
+      handles_mask[idx] = true;
+    }
+    for (int32 idx = 0; idx < node0->num_inputs(); ++idx) {
+      if (handles_mask[idx]) {
+        const Edge* edge0;
+        TF_RETURN_IF_ERROR(node0->input_edge(idx, &edge0));
+        packed_op_builder.Input(edge0->src());
+      } else {
+        std::vector<NodeBuilder::NodeOut> input_slices;
+        for (Node* node : c.second) {
+          const Edge* edge;
+          TF_RETURN_IF_ERROR(node->input_edge(idx, &edge));
+          input_slices.emplace_back(edge->src(), edge->src_output());
+        }
+        packed_op_builder.Input(input_slices);
+      }
+    }
+
+    Node* packed_op_node = nullptr;
+    TF_RETURN_IF_ERROR(packed_op_builder.Finalize(graph, &packed_op_node));
+    VLOG(1) << "Packed " << op_type_ << " ops (" << c.first << ") with "
+            << c.second.size() << " nodes on "
+            << packed_op_node->assigned_device_name() << ": "
+            << NodeJoin(c.second, ", ") << " -> " << packed_op_node->name()
+            << " in graph " << static_cast<void*>(graph);
+    occurrence_count++;
+
+    for (size_t i = 0; i < c.second.size(); ++i) {
+      std::vector<NodeBuilder::NodeOut> outputs;
+      outputs.reserve(c.second[i]->num_outputs());
+      for (int32 o = 0; o < c.second[i]->num_outputs(); ++o) {
+        outputs.emplace_back(packed_op_node, c.second.size() * o + i);
+      }
+      Node* output_op_node = nullptr;
+      TF_RETURN_IF_ERROR(
+          NodeBuilder(c.second[i]->name(), "IdentityN")
+              .Device(packed_op_node->requested_device())
+              .AssignedDevice(packed_op_node->assigned_device_name())
+              .Input(outputs)
+              .Finalize(graph, &output_op_node));
+
+      std::vector<const Edge*> edges;
+      for (const auto& edge : c.second[i]->out_edges()) {
+        edges.push_back(edge);
+      }
+      for (const auto& edge : edges) {
+        if (!edge) {
+          continue;
+        }
+        if (edge->IsControlEdge()) {
+          graph->AddControlEdge(output_op_node, edge->dst());
+        } else {
+          graph->AddEdge(output_op_node, edge->src_output(), edge->dst(),
+                         edge->dst_input());
+        }
+        graph->RemoveEdge(edge);
+      }
+
+      candidates.erase(c.second[i]);
+      graph->RemoveNode(c.second[i]);
+    }
+  }
+
+  if (poccurrence_count != nullptr) {
+    *poccurrence_count = occurrence_count;
+  }
+  return Status::OK();
+}
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/common/packing.h b/hybridbackend/tensorflow/graph/common/packing.h
new file mode 100644
index 00000000..eb681dae
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/packing.h
@@ -0,0 +1,68 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_PACKING_H_
+#define HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_PACKING_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include <tensorflow/core/graph/graph.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+class Pack {
+ public:
+  Pack(const string& op_type, const string& optimized_op_type);
+  Pack& WithDevice(const string& device);
+  Pack& WithTypeAttr(const string& attr_name,
+                     const std::vector<DataType>& constraints);
+  Pack& WithShapeAttr(const string& attr_name);
+  Pack& WithIntAttr(const string& attr_name);
+  Pack& WithStrAttr(const string& attr_name);
+  Pack& WithAggregatedShapeAttr(const string& attr_name);
+  Pack& WithAggregatedIntAttr(const string& attr_name);
+  Pack& WithAggregatedStrAttr(const string& attr_name);
+  Pack& WithHandle(const int32 input);
+  Pack& WithBuckets(const int32 num_buckets);
+  Status In(Graph* graph);
+  Status In(Graph* graph, int64* poccurrence_count);
+
+ private:
+  string op_type_;
+  string optimized_op_type_;
+  string device_;
+  std::map<string, std::vector<DataType>> type_attrs_;
+  std::vector<string> shape_attrs_;
+  std::vector<string> int_attrs_;
+  std::vector<string> str_attrs_;
+  std::vector<string> aggregated_shape_attrs_;
+  std::vector<string> aggregated_int_attrs_;
+  std::vector<string> aggregated_str_attrs_;
+  std::vector<int32> handles_;
+  int num_buckets_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(Pack);
+};
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
+#endif  // HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_PACKING_H_
diff --git a/hybridbackend/tensorflow/graph/common/relocation.cc b/hybridbackend/tensorflow/graph/common/relocation.cc
new file mode 100644
index 00000000..fdf76e12
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/relocation.cc
@@ -0,0 +1,199 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <string>
+#include <vector>
+
+#include <absl/strings/str_cat.h>
+#include <absl/strings/str_join.h>
+#include <tensorflow/core/graph/node_builder.h>
+#include <tensorflow/core/util/device_name_utils.h>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/helper.h"
+#include "hybridbackend/tensorflow/graph/common/relocation.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+RelocateOutputs::RelocateOutputs(const string& op_type_)
+    : op_type_(op_type_), device_(DEVICE_GPU) {}
+
+RelocateOutputs& RelocateOutputs::WithDevice(const string& device) {
+  device_ = device;
+  return *this;
+}
+
+RelocateOutputs& RelocateOutputs::Force() {
+  force_ = true;
+  return *this;
+}
+
+Status RelocateOutputs::In(Graph* graph) { return In(graph, nullptr); }
+
+Status RelocateOutputs::In(Graph* graph, int64* poccurrence_count) {
+  int64 occurrence_count = 0;
+  for (Node* node : graph->op_nodes()) {
+    if (!node->IsOp()) {
+      continue;
+    }
+
+    if (node->type_string() != op_type_) {
+      continue;
+    }
+
+    DeviceNameUtils::ParsedName root_device;
+    if (!DeviceNameUtils::ParseFullName(node->requested_device(),
+                                        &root_device)) {
+      continue;
+    }
+
+    DFSFrom(
+        *graph, {node},
+        [&](Node* n) {
+          if (!n->IsOp()) {
+            return;
+          }
+
+          string original_device = n->requested_device();
+          DeviceNameUtils::ParsedName requested_device;
+          if (!DeviceNameUtils::ParseFullName(original_device,
+                                              &requested_device)) {
+            return;
+          }
+
+          if (requested_device.job != root_device.job ||
+              requested_device.replica != root_device.replica ||
+              requested_device.task != root_device.task) {
+            VLOG(1) << "Skipped relocation of " << n->name()
+                    << " since task differs";
+            return;
+          }
+
+          if (requested_device.has_type) {
+            if (requested_device.type == device_) {
+              return;
+            }
+            if (!force_) {
+              VLOG(1) << "Skipped relocation of " << n->name()
+                      << " since device already set: " << requested_device.type;
+              return;
+            }
+          }
+
+          static constexpr char kResourceAttrName[] = "container";
+
+          if (HasNodeAttr(n->def(), kResourceAttrName)) {
+            VLOG(1) << "Skipped relocation of " << n->name()
+                    << " since resource in graph " << static_cast<void*>(graph);
+            return;
+          }
+
+          static constexpr char kColocationAttrName[] = "_class";
+          if (HasNodeAttr(n->def(), kColocationAttrName)) {
+            VLOG(1) << "Skipped relocation of " << n->name()
+                    << " since colocation in graph "
+                    << static_cast<void*>(graph);
+            return;
+          }
+
+          requested_device.type = device_;
+          requested_device.has_type = true;
+
+          auto relocated_device =
+              DeviceNameUtils::ParsedNameToString(requested_device);
+          n->set_requested_device(relocated_device);
+          VLOG(1) << "Relocated " << n->name() << " from \"" << original_device
+                  << "\" to \"" << relocated_device << "\" (" << device_
+                  << ") in graph " << static_cast<void*>(graph);
+          occurrence_count++;
+        },
+        nullptr);
+  }
+
+  if (poccurrence_count != nullptr) {
+    *poccurrence_count = occurrence_count;
+  }
+  return Status::OK();
+}
+
+Relocate::Relocate(const string& op_type_)
+    : op_type_(op_type_), device_(DEVICE_GPU), input_(0) {}
+
+Relocate& Relocate::WithDevice(const string& device) {
+  device_ = device;
+  return *this;
+}
+
+Relocate& Relocate::WithInput(const int32 input) {
+  input_ = input;
+  return *this;
+}
+
+Status Relocate::In(Graph* graph) { return In(graph, nullptr); }
+
+Status Relocate::In(Graph* graph, int64* poccurrence_count) {
+  int64 occurrence_count = 0;
+  for (Node* node : graph->op_nodes()) {
+    if (!node->IsOp()) {
+      continue;
+    }
+
+    if (node->type_string() != op_type_) {
+      continue;
+    }
+
+    const Edge* edge;
+    auto status = node->input_edge(input_, &edge);
+    if (!status.ok()) {
+      continue;
+    }
+
+    DeviceNameUtils::ParsedName requested_device;
+    if (!DeviceNameUtils::ParseFullName(node->requested_device(),
+                                        &requested_device)) {
+      continue;
+    }
+
+    if (requested_device.type == device_) {
+      continue;
+    }
+
+    static constexpr char kColocationAttrName[] = "_class";
+
+    if (node->requested_device() != edge->src()->requested_device()) {
+      if (HasNodeAttr(node->def(), kColocationAttrName)) {
+        node->ClearAttr(kColocationAttrName);
+      }
+      node->set_requested_device(edge->src()->requested_device());
+      VLOG(1) << "Relocated " << node->name() << " with input "
+              << edge->src()->name() << " in graph "
+              << static_cast<void*>(graph);
+      occurrence_count++;
+    }
+  }
+
+  if (poccurrence_count != nullptr) {
+    *poccurrence_count = occurrence_count;
+  }
+  return Status::OK();
+}
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/common/relocation.h b/hybridbackend/tensorflow/graph/common/relocation.h
new file mode 100644
index 00000000..49819823
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/relocation.h
@@ -0,0 +1,65 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_RELOCATION_H_
+#define HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_RELOCATION_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <string>
+#include <vector>
+
+#include <tensorflow/core/graph/graph.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+class RelocateOutputs {
+ public:
+  RelocateOutputs(const string& op_type);
+  RelocateOutputs& WithDevice(const string& device);
+  RelocateOutputs& Force();
+  Status In(Graph* graph);
+  Status In(Graph* graph, int64* poccurrence_count);
+
+ private:
+  string op_type_;
+  string device_;
+  bool force_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(RelocateOutputs);
+};
+
+class Relocate {
+ public:
+  Relocate(const string& op_type);
+  Relocate& WithDevice(const string& device);
+  Relocate& WithInput(const int32 input);
+  Status In(Graph* graph);
+  Status In(Graph* graph, int64* poccurrence_count);
+
+ private:
+  string op_type_;
+  string device_;
+  int32 input_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(Relocate);
+};
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
+#endif  // HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_RELOCATION_H_
diff --git a/hybridbackend/tensorflow/graph/common/replacing.cc b/hybridbackend/tensorflow/graph/common/replacing.cc
new file mode 100644
index 00000000..2d7e6a1a
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/replacing.cc
@@ -0,0 +1,190 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <algorithm>
+#include <deque>
+#include <unordered_map>
+#include <vector>
+
+#include <absl/strings/str_cat.h>
+#include <absl/strings/str_join.h>
+#include <absl/strings/str_split.h>
+
+#include <tensorflow/core/graph/node_builder.h>
+#include <tensorflow/core/util/device_name_utils.h>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/helper.h"
+#include "hybridbackend/tensorflow/graph/common/replacing.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+Replace::Replace(const string& op_type, const string& optimized_op_type)
+    : op_type_(op_type),
+      optimized_op_type_(optimized_op_type),
+      device_(DEVICE_GPU),
+      packed_(false) {}
+
+Replace& Replace::WithDevice(const string& device) {
+  device_ = device;
+  return *this;
+}
+
+Replace& Replace::WithTypeAttr(const string& attr_name,
+                               const std::vector<DataType>& constraints) {
+  type_attrs_[attr_name] = constraints;
+  return *this;
+}
+
+Replace& Replace::WithExtraIntAttr(const string& attr_name) {
+  extra_int_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Replace& Replace::Packed() {
+  packed_ = true;
+  return *this;
+}
+
+Status Replace::In(Graph* graph) { return In(graph, nullptr); }
+
+Status Replace::In(Graph* graph, int64* poccurrence_count) {
+  int64 occurrence_count = 0;
+  for (Node* node : graph->op_nodes()) {
+    if (!node->IsOp()) {
+      continue;
+    }
+
+    if (node->type_string() != op_type_) {
+      continue;
+    }
+
+    const int32 num_inputs = node->num_inputs();
+    if (num_inputs < 1) {
+      VLOG(2) << "Skipped replacing op " << node->name() << " since no inputs";
+      continue;
+    }
+
+    DeviceNameUtils::ParsedName requested_device;
+    if (!DeviceNameUtils::ParseFullName(node->requested_device(),
+                                        &requested_device)) {
+      VLOG(2) << "Skipped replacing op " << node->name()
+              << " since no device set";
+      continue;
+    }
+
+    if (requested_device.type != device_ || requested_device.job == "ps") {
+      VLOG(2) << "Skipped replacing op " << node->name()
+              << " since device differs, required /" << device_
+              << ":0 but requested: " << node->requested_device();
+      continue;
+    }
+
+    bool types_are_incompatible = false;
+    for (auto& attr_pair : type_attrs_) {
+      if (!HasNodeAttr(node->def(), attr_pair.first)) {
+        types_are_incompatible = true;
+        break;
+      }
+      DataType attrval;
+      TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attr_pair.first, &attrval));
+      if (attr_pair.second.size() == 0) {
+        continue;
+      }
+      types_are_incompatible = true;
+      for (DataType dt : attr_pair.second) {
+        if (dt == attrval) {
+          types_are_incompatible = false;
+          break;
+        }
+      }
+    }
+
+    if (TF_PREDICT_FALSE(types_are_incompatible)) {
+      VLOG(2) << "Skipped replacing op " << node->name()
+              << " since types incompatible";
+      continue;
+    }
+
+    auto name = node->name();
+    NodeBuilder replaced_op_builder = NodeBuilder(name, optimized_op_type_);
+    replaced_op_builder.Device(node->requested_device());
+    replaced_op_builder.AssignedDevice(node->assigned_device_name());
+    for (const auto& iter : node->attrs()) {
+      replaced_op_builder.Attr(iter.first, iter.second);
+    }
+
+    if (packed_) {
+      replaced_op_builder.Attr("N", 1);
+    }
+
+    std::vector<Node*> control_inputs;
+    for (const auto& edge : node->in_edges()) {
+      if (!edge) {
+        continue;
+      }
+      if (edge->IsControlEdge()) {
+        control_inputs.push_back(edge->src());
+      }
+    }
+    replaced_op_builder.ControlInputs(control_inputs);
+
+    for (int32 idx = 0; idx < node->num_inputs(); ++idx) {
+      const Edge* edge;
+      TF_RETURN_IF_ERROR(node->input_edge(idx, &edge));
+      if (packed_) {
+        replaced_op_builder.Input(gtl::ArraySlice<NodeBuilder::NodeOut>{
+            NodeBuilder::NodeOut(edge->src(), edge->src_output())});
+      } else {
+        replaced_op_builder.Input(
+            NodeBuilder::NodeOut(edge->src(), edge->src_output()));
+      }
+    }
+
+    Node* replaced_op_node = nullptr;
+    TF_RETURN_IF_ERROR(replaced_op_builder.Finalize(graph, &replaced_op_node));
+    VLOG(1) << "Replaced " << op_type_ << " op " << name << " on "
+            << node->requested_device() << " in graph "
+            << static_cast<void*>(graph);
+    occurrence_count++;
+
+    for (const auto& edge : node->out_edges()) {
+      if (!edge) {
+        continue;
+      }
+      if (edge->IsControlEdge()) {
+        graph->AddControlEdge(replaced_op_node, edge->dst());
+      } else {
+        graph->AddEdge(replaced_op_node, edge->src_output(), edge->dst(),
+                       edge->dst_input());
+      }
+    }
+
+    graph->RemoveNode(node);
+  }
+
+  if (poccurrence_count != nullptr) {
+    *poccurrence_count = occurrence_count;
+  }
+  return Status::OK();
+}
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/common/replacing.h b/hybridbackend/tensorflow/graph/common/replacing.h
new file mode 100644
index 00000000..e9bd1aa6
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/replacing.h
@@ -0,0 +1,56 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_REPLACING_H_
+#define HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_REPLACING_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include <tensorflow/core/graph/graph.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+class Replace {
+ public:
+  Replace(const string& op_type, const string& optimized_op_type);
+  Replace& WithDevice(const string& device);
+  Replace& WithTypeAttr(const string& attr_name,
+                        const std::vector<DataType>& constraints);
+  Replace& WithExtraIntAttr(const string& attr_name);
+  Replace& Packed();
+  Status In(Graph* graph);
+  Status In(Graph* graph, int64* poccurrence_count);
+
+ private:
+  string op_type_;
+  string optimized_op_type_;
+  string device_;
+  bool packed_;
+  std::map<string, std::vector<DataType>> type_attrs_;
+  std::vector<string> extra_int_attrs_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(Replace);
+};
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
+#endif  // HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_REPLACING_H_
diff --git a/hybridbackend/tensorflow/graph/common/rewriting.cc b/hybridbackend/tensorflow/graph/common/rewriting.cc
new file mode 100644
index 00000000..968e7b30
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/rewriting.cc
@@ -0,0 +1,247 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <algorithm>
+#include <deque>
+#include <unordered_map>
+#include <vector>
+
+#include <absl/strings/str_cat.h>
+
+#include <tensorflow/core/graph/node_builder.h>
+#include <tensorflow/core/util/device_name_utils.h>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/helper.h"
+#include "hybridbackend/tensorflow/graph/common/rewriting.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+Rewrite::Rewrite(const string& op_like_name, const string& op_name)
+    : op_like_name_(absl::StrCat("_HB_", op_like_name)),
+      op_name_(op_name),
+      device_(DEVICE_GPU) {}
+
+Rewrite& Rewrite::WithDevice(const string& device) {
+  device_ = device;
+  return *this;
+}
+
+Rewrite& Rewrite::WithTypeAttr(const string& attr_name,
+                               const DataType& default_attr) {
+  type_attrs_[attr_name] = default_attr;
+  return *this;
+}
+Rewrite& Rewrite::WithShapeAttr(const string& attr_name,
+                                const TensorShape& default_attr) {
+  shape_attrs_[attr_name] = default_attr;
+  return *this;
+}
+Rewrite& Rewrite::WithIntAttr(const string& attr_name,
+                              const int32& default_attr) {
+  int_attrs_[attr_name] = default_attr;
+  return *this;
+}
+Rewrite& Rewrite::WithStrAttr(const string& attr_name,
+                              const string& default_attr) {
+  str_attrs_[attr_name] = default_attr;
+  return *this;
+}
+
+Rewrite& Rewrite::WithTypeListAttr(const string& attr_name) {
+  type_list_attrs_.push_back(attr_name);
+  return *this;
+}
+
+Status Rewrite::In(Graph* graph) { return In(graph, nullptr); }
+
+Status Rewrite::In(Graph* graph, int64* poccurrence_count) {
+  int64 occurrence_count = 0;
+  for (Node* node : graph->op_nodes()) {
+    if (!node->IsOp()) {
+      continue;
+    }
+
+    // Output op should have op_like_name_ attribute.
+    if (!HasNodeAttr(node->def(), op_like_name_)) {
+      continue;
+    }
+
+    // Construct rewrited op.
+    NodeBuilder rewrited_op_builder = NodeBuilder(node->name(), op_name_);
+
+    // Place op to specific device.
+    if (!device_.empty()) {
+      DeviceNameUtils::ParsedName requested_device;
+      if (!DeviceNameUtils::ParseFullName(node->requested_device(),
+                                          &requested_device)) {
+        continue;
+      }
+      requested_device.type = device_;
+      rewrited_op_builder.Device(
+          DeviceNameUtils::ParsedNameToString(requested_device));
+    }
+
+    // String attributes.
+    for (auto& attr_pair : str_attrs_) {
+      const string attrkey = absl::StrCat(op_like_name_, "_", attr_pair.first);
+      string attrval = attr_pair.second;
+      if (HasNodeAttr(node->def(), attrkey)) {
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attrkey, &attrval));
+      }
+      rewrited_op_builder.Attr(attr_pair.first, attrval);
+    }
+
+    // Integer attributes.
+    for (auto& attr_pair : int_attrs_) {
+      const string attrkey = absl::StrCat(op_like_name_, "_", attr_pair.first);
+      int32 attrval = attr_pair.second;
+      if (HasNodeAttr(node->def(), attrkey)) {
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attrkey, &attrval));
+      }
+      rewrited_op_builder.Attr(attr_pair.first, attrval);
+    }
+
+    // Data type attributes.
+    for (auto& attr_pair : type_attrs_) {
+      const string attrkey = absl::StrCat(op_like_name_, "_", attr_pair.first);
+      DataType attrval = attr_pair.second;
+      if (HasNodeAttr(node->def(), attrkey)) {
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attrkey, &attrval));
+      }
+      rewrited_op_builder.Attr(attr_pair.first, attrval);
+    }
+
+    // Tensor shape attributes.
+    for (auto& attr_pair : shape_attrs_) {
+      const string attrkey = absl::StrCat(op_like_name_, "_", attr_pair.first);
+      TensorShape attrval = attr_pair.second;
+      if (HasNodeAttr(node->def(), attrkey)) {
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attrkey, &attrval));
+      }
+      rewrited_op_builder.Attr(attr_pair.first, attrval);
+    }
+
+    // Data type list attributes.
+    for (auto& attr : type_list_attrs_) {
+      const string attrkey = absl::StrCat(op_like_name_, "_", attr);
+      std::vector<DataType> attrval;
+      if (HasNodeAttr(node->def(), attrkey)) {
+        TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attrkey, &attrval));
+      }
+      rewrited_op_builder.Attr(attr, attrval);
+    }
+
+    std::vector<Node*> control_inputs;
+    const string input_proxy_key = absl::StrCat(op_like_name_, "_input_proxy");
+    std::unordered_map<Node*, int> input_proxies;
+    for (const auto& edge : node->in_edges()) {
+      if (!edge) {
+        continue;
+      }
+      if (!edge->IsControlEdge()) {
+        continue;
+      }
+      Node* control_input = edge->src();
+      if (HasNodeAttr(control_input->def(), input_proxy_key)) {
+        int32 input_idx = -1;
+        TF_RETURN_IF_ERROR(
+            GetNodeAttr(control_input->attrs(), input_proxy_key, &input_idx));
+        if (input_idx < 0) {
+          LOG(ERROR) << "Invalid attribute set for input proxy: "
+                     << control_input->name();
+          control_inputs.push_back(control_input);
+        } else {
+          input_proxies[control_input] = input_idx;
+        }
+      } else {
+        control_inputs.push_back(control_input);
+      }
+    }
+
+    std::vector<Node*> input_proxy_vec(input_proxies.size(), nullptr);
+    for (auto& it : input_proxies) {
+      if (it.second < 0 || it.second > input_proxy_vec.size() - 1) {
+        return errors::Internal("Invalid attribute set for input proxy: ",
+                                it.first->name());
+      }
+      input_proxy_vec[it.second] = it.first;
+    }
+    for (size_t idx = 0; idx < input_proxy_vec.size(); ++idx) {
+      if (input_proxy_vec[idx] == nullptr) {
+        return errors::Internal("Invalid attribute set for input proxy ", idx);
+      }
+    }
+
+    for (Node* input_proxy : input_proxy_vec) {
+      if (input_proxy->type_string() == "Identity") {
+        const Edge* edge;
+        TF_RETURN_IF_ERROR(input_proxy->input_edge(0, &edge));
+        rewrited_op_builder.Input(
+            NodeBuilder::NodeOut(edge->src(), edge->src_output()));
+      } else if (input_proxy->type_string() == "IdentityN") {
+        std::vector<NodeBuilder::NodeOut> list_inputs;
+        list_inputs.reserve(input_proxy->num_inputs());
+        for (int32 idx = 0; idx < input_proxy->num_inputs(); ++idx) {
+          const Edge* edge;
+          TF_RETURN_IF_ERROR(input_proxy->input_edge(idx, &edge));
+          list_inputs.emplace_back(edge->src(), edge->src_output());
+        }
+        rewrited_op_builder.Input(list_inputs);
+      } else {
+        return errors::Internal("Invalid type ", input_proxy->type_string(),
+                                " for input proxy ", input_proxy->name());
+      }
+    }
+    rewrited_op_builder.ControlInputs(control_inputs);
+
+    Node* rewrited_op_node = nullptr;
+    TF_RETURN_IF_ERROR(rewrited_op_builder.Finalize(graph, &rewrited_op_node));
+    occurrence_count++;
+    VLOG(1) << "Rewrited function " << op_like_name_ << " to " << op_name_
+            << " op " << rewrited_op_node->name() << " in graph "
+            << static_cast<void*>(graph);
+
+    for (const auto& edge : node->out_edges()) {
+      if (!edge) {
+        continue;
+      }
+      if (edge->IsControlEdge()) {
+        graph->AddControlEdge(rewrited_op_node, edge->dst());
+      } else {
+        graph->AddEdge(rewrited_op_node, edge->src_output(), edge->dst(),
+                       edge->dst_input());
+      }
+    }
+
+    graph->RemoveNode(node);
+    for (auto& it : input_proxies) {
+      graph->RemoveNode(it.first);
+    }
+  }
+
+  if (poccurrence_count != nullptr) {
+    *poccurrence_count = occurrence_count;
+  }
+  return Status::OK();
+}
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/common/rewriting.h b/hybridbackend/tensorflow/graph/common/rewriting.h
new file mode 100644
index 00000000..8a58e3de
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/common/rewriting.h
@@ -0,0 +1,60 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_REWRITING_H_
+#define HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_REWRITING_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <string>
+#include <vector>
+
+#include <tensorflow/core/graph/graph.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+class Rewrite {
+ public:
+  Rewrite(const string& op_like_name, const string& op_name);
+  Rewrite& WithDevice(const string& device);
+  Rewrite& WithTypeAttr(const string& attr_name, const DataType& default_attr);
+  Rewrite& WithShapeAttr(const string& attr_name,
+                         const TensorShape& default_attr);
+  Rewrite& WithIntAttr(const string& attr_name, const int32& default_attr);
+  Rewrite& WithStrAttr(const string& attr_name, const string& default_attr);
+  Rewrite& WithTypeListAttr(const string& attr_name);
+  Status In(Graph* graph);
+  Status In(Graph* graph, int64* poccurrence_count);
+
+ private:
+  string op_like_name_;
+  string op_name_;
+  string device_;
+  int32 num_inputs_;
+  std::map<string, DataType> type_attrs_;
+  std::map<string, TensorShape> shape_attrs_;
+  std::map<string, int32> int_attrs_;
+  std::map<string, string> str_attrs_;
+  std::vector<string> type_list_attrs_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(Rewrite);
+};
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
+#endif  // HYBRIDBACKEND_TENSORFLOW_GRAPH_COMMON_REWRITING_H_
diff --git a/hybridbackend/tensorflow/graph/op_optimization.cc b/hybridbackend/tensorflow/graph/op_optimization.cc
new file mode 100644
index 00000000..c322d978
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/op_optimization.cc
@@ -0,0 +1,140 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <unistd.h>
+#include <fstream>
+#include <vector>
+
+#include <absl/strings/str_cat.h>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/relocation.h"
+#include "hybridbackend/tensorflow/graph/op_optimization.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+namespace {
+inline bool OpOptimizationDisabled() {
+  static const bool kOptimizationDisabled =
+      ::hybridbackend::EnvVarGetBool("HB_OP_OPTIMIZATION_DISABLED", false);
+  return kOptimizationDisabled;
+}
+
+inline bool OpRelocationEnabled() {
+  static const bool kRelocationEnabled =
+      ::hybridbackend::EnvVarGetBool("HB_OP_RELOCATION_ENABLED", false);
+  return kRelocationEnabled;
+}
+}  // namespace
+
+Status OpOptimizationPass::Run(const GraphOptimizationPassOptions& options) {
+  if (options.graph == nullptr) {
+    return Status::OK();
+  }
+
+  Graph* graph = options.graph->get();
+  if (graph == nullptr) {
+    return errors::Internal(
+        "Op optimization should happen before partitioning "
+        "and a graph should be available.");
+  }
+
+  bool opt_disabled = false;
+
+#if HYBRIDBACKEND_CHECK_INSTANCE
+  static long kTwoSeconds = 2000L;
+  static bool kInstanceChecked = ::hybridbackend::EnvCheckInstance(kTwoSeconds);
+  if (TF_PREDICT_FALSE(!kInstanceChecked)) {
+    opt_disabled = true;
+  }
+#endif
+
+  opt_disabled |= OpOptimizationDisabled();
+
+  TF_RETURN_IF_ERROR(Optimize(graph, options.session_options, opt_disabled));
+  return Status::OK();
+}
+
+class OpOptimizationBeginPass : public GraphOptimizationPass {
+ public:
+  Status Run(const GraphOptimizationPassOptions& options) override {
+    if (options.graph == nullptr) {
+      return Status::OK();
+    }
+
+    Graph* graph = options.graph->get();
+    if (graph == nullptr) {
+      return errors::Internal(
+          "Op optimization should happen before partitioning "
+          "and a graph should be available.");
+    }
+
+    static const string kUnoptimizedGraphPath =
+        ::hybridbackend::EnvVarGet("HB_UNOPTIMIZED_GRAPH_PATH", "");
+
+    if (!kUnoptimizedGraphPath.empty()) {
+      std::ofstream ofs(absl::StrCat(kUnoptimizedGraphPath, ".", getpid(),
+                                     "_0x", absl::Hex(graph), ".pbtxt"));
+      ofs << graph->ToGraphDefDebug().DebugString();
+    }
+
+    if (TF_PREDICT_FALSE(OpRelocationEnabled())) {
+      RelocateOutputs("Unique").Force().In(graph);
+    }
+
+    return Status::OK();
+  }
+};
+
+REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 0,
+                      OpOptimizationBeginPass);
+
+class OpOptimizationEndPass : public GraphOptimizationPass {
+ public:
+  Status Run(const GraphOptimizationPassOptions& options) override {
+    if (options.graph == nullptr) {
+      return Status::OK();
+    }
+
+    Graph* graph = options.graph->get();
+    if (graph == nullptr) {
+      return errors::Internal(
+          "Op optimization should happen before partitioning "
+          "and a graph should be available.");
+    }
+
+    static const string kOptimizedGraphPath =
+        ::hybridbackend::EnvVarGet("HB_OPTIMIZED_GRAPH_PATH", "");
+
+    if (!kOptimizedGraphPath.empty()) {
+      std::ofstream ofs(absl::StrCat(kOptimizedGraphPath, ".", getpid(), "_0x",
+                                     absl::Hex(graph), ".pbtxt"));
+      ofs << graph->ToGraphDefDebug().DebugString();
+    }
+
+    return Status::OK();
+  }
+};
+
+REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 1000,
+                      OpOptimizationEndPass);
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/op_optimization.h b/hybridbackend/tensorflow/graph/op_optimization.h
new file mode 100644
index 00000000..b689305b
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/op_optimization.h
@@ -0,0 +1,49 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef HYBRIDBACKEND_TENSORFLOW_GRAPH_OP_OPTIMIZATION_H_
+#define HYBRIDBACKEND_TENSORFLOW_GRAPH_OP_OPTIMIZATION_H_
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <string>
+#include <vector>
+
+#include <tensorflow/core/common_runtime/optimization_registry.h>
+
+namespace tensorflow {
+namespace hybridbackend {
+
+class OpOptimizationPass : public GraphOptimizationPass {
+ public:
+  virtual Status Run(const GraphOptimizationPassOptions& options);
+
+ protected:
+  virtual Status Optimize(Graph* graph, const SessionOptions* options,
+                          const bool disabled) = 0;
+};
+
+#define REGISTER_REPLACING_OPTIMIZATION(PASS) \
+  REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 1, PASS)
+
+#define REGISTER_REDUCTION_OPTIMIZATION(PASS)                                 \
+  REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 100, \
+                        PASS)
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
+#endif  // HYBRIDBACKEND_TENSORFLOW_GRAPH_OP_OPTIMIZATION_H_
diff --git a/hybridbackend/tensorflow/graph/optimize_collective.cc b/hybridbackend/tensorflow/graph/optimize_collective.cc
new file mode 100644
index 00000000..a01032f5
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/optimize_collective.cc
@@ -0,0 +1,168 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <vector>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/linearization.h"
+#include "hybridbackend/tensorflow/graph/common/packing.h"
+#include "hybridbackend/tensorflow/graph/common/rewriting.h"
+#include "hybridbackend/tensorflow/graph/op_optimization.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+namespace {
+inline bool CollectivePackingDisabled() {
+  static const bool kCollectivePackingDisabled =
+      ::hybridbackend::EnvVarGetBool("HB_COLLECTIVE_PACKING_DISABLED", false);
+  return kCollectivePackingDisabled;
+}
+inline string CollectivePackingAlgorithm() {
+  // MERGE or GROUP allowed
+  static const string kCollectivePackingAlgorithm =
+      ::hybridbackend::EnvVarGet("HB_COLLECTIVE_PACKING_ALGORITHM", "GROUP");
+  return kCollectivePackingAlgorithm;
+}
+}  // namespace
+
+class OptimizeCollectiveReplacingPass : public OpOptimizationPass {
+ public:
+  Status Optimize(Graph* graph, const SessionOptions* options,
+                  const bool disabled) override {
+    TF_RETURN_IF_ERROR(Rewrite("GetCollectiveId", "HbGetNcclId").In(graph));
+    TF_RETURN_IF_ERROR(Rewrite("CollectiveHandleOp", "HbNcclCollectiveHandleOp")
+                           .WithStrAttr("container", "")
+                           .WithStrAttr("shared_name", "")
+                           .In(graph));
+    TF_RETURN_IF_ERROR(Rewrite("CreateCollective", "HbCreateNcclCollective")
+                           .WithStrAttr("shared_name", "")
+                           .WithIntAttr("world_size", 1)
+                           .WithIntAttr("local_size", 1)
+                           .WithIntAttr("rank", 0)
+                           .In(graph));
+    TF_RETURN_IF_ERROR(
+        Rewrite("IsCollectiveInitialized", "HbIsNcclCollectiveInitialized")
+            .In(graph));
+
+    // Attribute T should be infered not specified.
+    TF_RETURN_IF_ERROR(Rewrite("CollectiveAllreduce", "HbNcclAllreduce")
+                           .WithIntAttr("reduce_op", 0)
+                           .In(graph));
+    // Attribute T should be infered not specified.
+    TF_RETURN_IF_ERROR(Rewrite("CollectiveAlltoall", "HbNcclAlltoall")
+                           .WithIntAttr("topology", 0)
+                           .WithTypeAttr("wire_dtype", DT_FLOAT)
+                           .In(graph));
+    // Attribute T should be infered not specified.
+    TF_RETURN_IF_ERROR(Rewrite("CollectiveAlltoallv", "HbNcclAlltoallv")
+                           .WithIntAttr("topology", 0)
+                           .WithTypeAttr("wire_dtype", DT_FLOAT)
+                           .WithShapeAttr("common_shape", TensorShape({}))
+                           .In(graph));
+    // Attribute T should be infered not specified.
+    TF_RETURN_IF_ERROR(Rewrite("CollectiveBroadcast", "HbNcclBroadcast")
+                           .WithIntAttr("root_rank", 0)
+                           .In(graph));
+    // Attribute T should be infered not specified.
+    TF_RETURN_IF_ERROR(
+        Rewrite("CollectiveAllgather", "HbNcclAllgather").In(graph));
+    // Attribute T should be infered not specified.
+    TF_RETURN_IF_ERROR(
+        Rewrite("CollectiveAllgatherv", "HbNcclAllgatherv").In(graph));
+
+    return Status::OK();
+  }
+};
+
+REGISTER_REPLACING_OPTIMIZATION(OptimizeCollectiveReplacingPass);
+
+class OptimizeCollectiveReductionPass : public OpOptimizationPass {
+ public:
+  Status Optimize(Graph* graph, const SessionOptions* options,
+                  const bool disabled) override {
+    if (TF_PREDICT_TRUE(!disabled && !CollectivePackingDisabled())) {
+      const std::vector<DataType> nccl_dtypes{DT_FLOAT, DT_DOUBLE, DT_HALF,
+                                              DT_INT8,  DT_INT32,  DT_INT64,
+                                              DT_UINT8, DT_UINT32, DT_INT64};
+      const std::vector<DataType> nccl_wire_dtypes{DT_FLOAT, DT_HALF};
+      if (TF_PREDICT_FALSE(CollectivePackingAlgorithm() == "GROUP")) {
+        TF_RETURN_IF_ERROR(Pack("HbNcclAllreduce", "HbNcclAllreduceN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithIntAttr("reduce_op")
+                               .WithHandle(0)
+                               .In(graph));
+        TF_RETURN_IF_ERROR(Pack("HbNcclAlltoall", "HbNcclAlltoallN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithTypeAttr("wire_dtype", nccl_wire_dtypes)
+                               .WithHandle(0)
+                               .In(graph));
+        TF_RETURN_IF_ERROR(Pack("HbNcclAlltoallv", "HbNcclAlltoallvN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithTypeAttr("wire_dtype", nccl_wire_dtypes)
+                               .WithAggregatedShapeAttr("common_shape")
+                               .WithHandle(0)
+                               .In(graph));
+      } else if (TF_PREDICT_FALSE(CollectivePackingAlgorithm() == "MERGE")) {
+        TF_RETURN_IF_ERROR(Pack("HbNcclAllreduce", "HbNcclAllreduceMergedN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithIntAttr("reduce_op")
+                               .WithHandle(0)
+                               .In(graph));
+        TF_RETURN_IF_ERROR(Pack("HbNcclAlltoall", "HbNcclAlltoallMergedN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithTypeAttr("wire_dtype", nccl_wire_dtypes)
+                               .WithHandle(0)
+                               .In(graph));
+        TF_RETURN_IF_ERROR(Pack("HbNcclAlltoallv", "HbNcclAlltoallvMergedN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithTypeAttr("wire_dtype", nccl_wire_dtypes)
+                               .WithShapeAttr("common_shape")
+                               .WithHandle(0)
+                               .In(graph));
+      } else {
+        TF_RETURN_IF_ERROR(Pack("HbNcclAllreduce", "HbNcclAllreduceMergedN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithIntAttr("reduce_op")
+                               .WithHandle(0)
+                               .In(graph));
+        TF_RETURN_IF_ERROR(Pack("HbNcclAlltoall", "HbNcclAlltoallMergedN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithTypeAttr("wire_dtype", nccl_wire_dtypes)
+                               .WithHandle(0)
+                               .In(graph));
+        TF_RETURN_IF_ERROR(Pack("HbNcclAlltoallv", "HbNcclAlltoallvN")
+                               .WithTypeAttr("dtype", nccl_dtypes)
+                               .WithTypeAttr("wire_dtype", nccl_wire_dtypes)
+                               .WithAggregatedShapeAttr("common_shape")
+                               .WithHandle(0)
+                               .In(graph));
+      }
+    }
+
+    TF_RETURN_IF_ERROR(
+        LinearizeOutputs("HbNcclCollectiveHandleOp", 0).In(graph));
+    return Status::OK();
+  }
+};
+
+REGISTER_REDUCTION_OPTIMIZATION(OptimizeCollectiveReductionPass);
+
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/graph/optimize_partition_by_modulo.cc b/hybridbackend/tensorflow/graph/optimize_partition_by_modulo.cc
new file mode 100644
index 00000000..fdb7484a
--- /dev/null
+++ b/hybridbackend/tensorflow/graph/optimize_partition_by_modulo.cc
@@ -0,0 +1,91 @@
+/* Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if HYBRIDBACKEND_TENSORFLOW
+
+#include <vector>
+
+#include "hybridbackend/common/env.h"
+#include "hybridbackend/tensorflow/graph/common/packing.h"
+#include "hybridbackend/tensorflow/graph/common/relocation.h"
+#include "hybridbackend/tensorflow/graph/common/replacing.h"
+#include "hybridbackend/tensorflow/graph/common/rewriting.h"
+#include "hybridbackend/tensorflow/graph/op_optimization.h"
+
+namespace tensorflow {
+namespace hybridbackend {
+
+namespace {
+inline bool PartitionByModuloOptimizationDisabled() {
+  static const bool kPartitionByModuloOptimizationDisabled =
+      ::hybridbackend::EnvVarGetBool(
+          "HB_OP_PARTITION_BY_MODULO_OPTIMIZATION_DISABLED", false);
+  return kPartitionByModuloOptimizationDisabled;
+}
+
+inline bool PartitionByModuloPackingDisabled() {
+  static const bool kPartitionByModuloPackingDisabled =
+      ::hybridbackend::EnvVarGetBool(
+          "HB_OP_PARTITION_BY_MODULO_PACKING_DISABLED", false);
+  return kPartitionByModuloPackingDisabled;
+}
+
+}  // namespace
+
+class OptimizePartitionByModuloReplacingPass : public OpOptimizationPass {
+ public:
+  Status Optimize(Graph* graph, const SessionOptions* options,
+                  const bool disabled) override {
+    if (TF_PREDICT_FALSE(disabled || PartitionByModuloOptimizationDisabled())) {
+      return Status::OK();
+    }
+
+    TF_RETURN_IF_ERROR(Rewrite("PartitionByModulo", "HbPartitionByModulo")
+                           .WithIntAttr("num_partitions", 1)
+                           .In(graph));
+
+    return Status::OK();
+  }
+};
+
+REGISTER_REPLACING_OPTIMIZATION(OptimizePartitionByModuloReplacingPass);
+
+class OptimizePartitionByModuloReductionPass : public OpOptimizationPass {
+ public:
+  Status Optimize(Graph* graph, const SessionOptions* options,
+                  const bool disabled) override {
+    if (TF_PREDICT_FALSE(disabled || PartitionByModuloOptimizationDisabled())) {
+      return Status::OK();
+    }
+
+    if (TF_PREDICT_TRUE(!PartitionByModuloPackingDisabled())) {
+      TF_RETURN_IF_ERROR(
+          Pack("HbPartitionByModulo", "HbPartitionByModuloN")
+              .WithTypeAttr("T", {DT_INT32, DT_INT64, DT_UINT32, DT_UINT64})
+              .WithIntAttr("num_partitions")
+              .In(graph));
+
+      return Status::OK();
+    }
+
+    return Status::OK();
+  }
+};
+
+REGISTER_REDUCTION_OPTIMIZATION(OptimizePartitionByModuloReductionPass);
+}  // namespace hybridbackend
+}  // namespace tensorflow
+
+#endif  // HYBRIDBACKEND_TENSORFLOW
diff --git a/hybridbackend/tensorflow/keras/layers/__init__.py b/hybridbackend/tensorflow/keras/layers/__init__.py
index fdfcdbae..0d69e859 100644
--- a/hybridbackend/tensorflow/keras/layers/__init__.py
+++ b/hybridbackend/tensorflow/keras/layers/__init__.py
@@ -20,5 +20,27 @@
 from __future__ import division
 from __future__ import print_function
 
-from hybridbackend.tensorflow.feature_column.dense_features import \
-  DenseFeatures
+try:
+  from tensorflow.python.feature_column.dense_features import DenseFeatures
+except ImportError:
+  pass
+
+
+def dense_features(features, feature_columns):
+  r'''Function produces dense tensors based on given `feature_columns`.
+
+  Args:
+    features: A mapping from key to tensors. `FeatureColumn`s look up via
+      these keys. For example `numeric_column('price')` will look at 'price'
+      key in this dict. Values can be a `SparseTensor` or a `Tensor` depends
+      on corresponding `FeatureColumn`.
+    feature_columns: List of feature columns.
+
+  Returns:
+    List of `Tensor`s which represents input layer of a model, which matches
+    order of columns in `feature_columns`.
+  '''
+  cols_to_output_tensors = {}
+  DenseFeatures(feature_columns)(
+    features, cols_to_output_tensors=cols_to_output_tensors)
+  return [cols_to_output_tensors[f] for f in feature_columns]
diff --git a/hybridbackend/tensorflow/keras/model.py b/hybridbackend/tensorflow/keras/model.py
index d9f099b6..10132dc8 100644
--- a/hybridbackend/tensorflow/keras/model.py
+++ b/hybridbackend/tensorflow/keras/model.py
@@ -41,9 +41,15 @@
 from tensorflow.python.keras.engine import training as _keras_training
 from tensorflow.python.keras.engine import training_arrays as _training_arrays
 from tensorflow.python.keras.engine import training_utils as _training_utils
+from tensorflow.python.keras.optimizers import TFOptimizer as _tf_optimizer
 from tensorflow.python.keras.utils import tf_utils
-from tensorflow.python.keras.utils.mode_keys import ModeKeys as _ModeKeys
+
+try:
+  from tensorflow.python.keras.utils.mode_keys import ModeKeys as _ModeKeys
+except ImportError:
+  pass
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import clip_ops
 from tensorflow.python.ops import gen_resource_variable_ops
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import state_ops
@@ -125,6 +131,86 @@ def __exit__(self, exc_type, exc_val, exc_tb):
       PatchGetSessionForKerasModel._stack_depth -= 1
 
 
+class PatchOptimizerForKerasModel(object):  # pylint: disable=useless-object-inheritance
+  r'''Context manager that patches TF APIs for `tf.keras.Optimizers`.
+  '''
+  _lock = threading.Lock()
+  _stack_depth = 0
+
+  def __init__(
+      self, clipnorm, clipvalue):
+    self._clipnorm = clipnorm
+    self._clipvalue = clipvalue
+
+  def __enter__(self):
+    with PatchOptimizerForKerasModel._lock:
+      PatchOptimizerForKerasModel._stack_depth += 1
+      if PatchOptimizerForKerasModel._stack_depth <= 1:
+
+        def wraps_get_updates():  # pylint: disable=unused-argument
+          r'''replace the default TFOptimizer's get_gradients.
+          '''
+
+          def wrapped_get_updates(cls, loss, params):  # pylint: disable=unused-argument
+            if distribution_strategy_context.has_strategy():
+              cls.updates = []
+
+              if not params:
+                grads = cls.optimizer.compute_gradients(loss)
+              else:
+                grads = cls.optimizer.compute_gradients(loss, params)
+              if self._clipnorm is not None:
+                grads = [
+                  (clip_ops.clip_by_norm(g[0], self._clipnorm)
+                   if g[0] is not None else g[0], g[1])
+                  for g in grads
+                ]
+              if self._clipvalue is not None:
+                grads = [
+                  (clip_ops.clip_by_value(
+                    g[0], -self._clipvalue, self._clipvalue)
+                    if g[0] is not None else g[0], g[1])
+                  for g in grads
+                ]
+              global_step = training_util.get_global_step()
+              opt_update = cls.optimizer.apply_gradients(grads, global_step)
+            else:
+              if not params:
+                cls.updates = [state_ops.assign_add(cls.iterations, 1)]
+                return cls.updates
+
+              cls.updates = []
+              grads = cls.optimizer.compute_gradients(loss, params)
+              if self._clipnorm is not None:
+                grads = [
+                  (clip_ops.clip_by_norm(g[0], self._clipnorm)
+                   if g[0] is not None else g[0], g[1])
+                  for g in grads
+                ]
+              if self._clipvalue is not None:
+                grads = [
+                  (clip_ops.clip_by_value(
+                    g[0], -self._clipvalue, self._clipvalue)
+                    if g[0] is not None else g[0], g[1])
+                  for g in grads
+                ]
+              opt_update = cls.optimizer.apply_gradients(
+                grads, global_step=cls.iterations)
+            cls.updates.append(opt_update)
+            return cls.updates
+          return wrapped_get_updates
+        self._prev_get_updates = _tf_optimizer.get_updates
+        _tf_optimizer.get_updates = wraps_get_updates()
+
+      return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    with PatchOptimizerForKerasModel._lock:
+      if PatchOptimizerForKerasModel._stack_depth <= 1:
+        _tf_optimizer.get_updates = self._prev_get_updates
+      PatchOptimizerForKerasModel._stack_depth -= 1
+
+
 class PatchCallbacksForKerasModel(object):  # pylint: disable=useless-object-inheritance
   r'''Context manager that patches CallbackList of `keras.Model`.
   '''
@@ -469,6 +555,12 @@ def compile(self,
       self._run_eagerly = kwargs.pop('run_eagerly', None)
       self._experimental_run_tf_function = kwargs.pop(
         'experimental_run_tf_function', True)
+      self._clipnorm = kwargs.pop('clipnorm', None)
+      if self._clipnorm is not None and self._clipnorm < 0:
+        raise ValueError(f'Expected clipnorm >=0, received: {self._clipnorm}')
+      self._clipvalue = kwargs.pop('clipvalue', None)
+      if self._clipvalue is not None and self._clipvalue < 0:
+        raise ValueError(f'Expected clipvalue >=0, received: {self._clipvalue}')
       self._set_optimizer(optimizer)
       is_any_optimizer_v1 = any(isinstance(opt, optimizers.Optimizer)
                                 for opt in nest.flatten(self.optimizer))
@@ -612,8 +704,9 @@ def fit(self, *args, **kwargs):
               y = labels
 
           validation_data = kwargs.pop('validation_data', None)
-          if validation_data is not None and isinstance(
-              validation_data, dataset_ops.DatasetV2):
+          if validation_data is not None:
+            if not isinstance(validation_data, dataset_ops.DatasetV2):
+              raise ValueError('validation_data must be a Dataset')
             with Context.scope(mode=ModeKeys.EVAL):
               val_feature, val_labels = dataset_ops.make_one_shot_iterator(
                 validation_data.repeat()).get_next()
@@ -825,15 +918,16 @@ def _make_train_function(self):
           K.get_graph()._unsafe_unfinalize()  # pylint: disable=protected-access
         with K.get_graph().as_default() as g, g.device(self._device_fn):
           with K.name_scope('training'):
-            # Training updates
-            updates = self.optimizer.get_updates(
-              params=self._collected_trainable_weights, loss=self.total_loss)
-            # Unconditional updates
-            updates += self.get_updates_for(None)
-            # Conditional updates relevant to this model
-            updates += self.get_updates_for(self.inputs)
-            global_step = training_util.get_or_create_global_step()
-            updates += [state_ops.assign_add(global_step, 1)]
+            with PatchOptimizerForKerasModel(self._clipnorm, self._clipvalue):
+              # Training updates
+              updates = self.optimizer.get_updates(
+                params=self._collected_trainable_weights, loss=self.total_loss)
+              # Unconditional updates
+              updates += self.get_updates_for(None)
+              # Conditional updates relevant to this model
+              updates += self.get_updates_for(self.inputs)
+              global_step = training_util.get_or_create_global_step()
+              updates += [state_ops.assign_add(global_step, 1)]
 
           metrics = self._get_training_eval_metrics()
           metrics_tensors = [
diff --git a/hybridbackend/tensorflow/metrics/auc.py b/hybridbackend/tensorflow/metrics/auc.py
index cd1d2250..2c031bf2 100644
--- a/hybridbackend/tensorflow/metrics/auc.py
+++ b/hybridbackend/tensorflow/metrics/auc.py
@@ -31,9 +31,8 @@
 from tensorflow.python.ops import weights_broadcast_ops
 from tensorflow.python.platform import tf_logging as logging
 
-from hybridbackend.tensorflow.distribute.communicator import CollectiveOps
-from hybridbackend.tensorflow.distribute.communicator_pool import \
-  CommunicatorPool
+from hybridbackend.tensorflow.distribute.collective import Collective
+from hybridbackend.tensorflow.distribute.ops import CollectiveOps
 
 
 def _confusion_matrix_at_thresholds(labels,
@@ -167,8 +166,8 @@ def _confusion_matrix_at_thresholds(labels,
   stacked = array_ops.stack([tp_sum, fn_sum, tn_sum, fp_sum])
   if isinstance(stacked, (list, tuple)):
     stacked = stacked[0]
-  sum_stacked = CommunicatorPool.get().allreduce(
-    stacked, reduce_op=CollectiveOps.SUM, trainable=False)
+  sum_stacked = Collective.get().allreduce(
+    stacked, reduce_op=CollectiveOps.SUM)
   if isinstance(sum_stacked, (list, tuple)):
     sum_stacked = sum_stacked[0]
   tp_sum, fn_sum, tn_sum, fp_sum = array_ops.unstack(sum_stacked)
diff --git a/hybridbackend/tensorflow/metrics/mean.py b/hybridbackend/tensorflow/metrics/mean.py
index 32eccd71..b625cf68 100644
--- a/hybridbackend/tensorflow/metrics/mean.py
+++ b/hybridbackend/tensorflow/metrics/mean.py
@@ -29,9 +29,8 @@
 from tensorflow.python.ops import variable_scope as vs
 from tensorflow.python.ops import weights_broadcast_ops
 
-from hybridbackend.tensorflow.distribute.communicator import CollectiveOps
-from hybridbackend.tensorflow.distribute.communicator_pool import \
-  CommunicatorPool
+from hybridbackend.tensorflow.distribute.collective import Collective
+from hybridbackend.tensorflow.distribute.ops import CollectiveOps
 
 
 def mean(values,
@@ -97,8 +96,8 @@ def mean(values,
     stacked = array_ops.stack([values_sum, num_values])
     if isinstance(stacked, (list, tuple)):
       stacked = stacked[0]
-    sum_stacked = CommunicatorPool.get().allreduce(
-      stacked, reduce_op=CollectiveOps.SUM, trainable=False)
+    sum_stacked = Collective.get().allreduce(
+      stacked, reduce_op=CollectiveOps.SUM)
     if isinstance(sum_stacked, (list, tuple)):
       sum_stacked = sum_stacked[0]
     values_sum, num_values = array_ops.unstack(sum_stacked)
diff --git a/hybridbackend/tensorflow/ops/__init__.py b/hybridbackend/tensorflow/ops/__init__.py
index e3866870..169853b4 100644
--- a/hybridbackend/tensorflow/ops/__init__.py
+++ b/hybridbackend/tensorflow/ops/__init__.py
@@ -20,8 +20,34 @@
 from __future__ import division
 from __future__ import print_function
 
-from hybridbackend.tensorflow.ops.floormod_shuffle.ops import \
-  floormod_partition
-from hybridbackend.tensorflow.ops.floormod_shuffle.ops import floormod_shuffle
-from hybridbackend.tensorflow.ops.floormod_shuffle.ops import \
-  floormod_shuffle_n
+import os as _os
+
+
+def disable_optimization():
+  r'''Disable optimizations for operators on GPU.
+  '''
+  _os.environ['HB_OP_OPTIMIZATION_DISABLED'] = '1'
+
+
+def enable_optimization(logging_level=None, relocate_ops=False):
+  r'''Enable optimizations for operators on GPU.
+
+  Args:
+    logging_level: Level of details to optimize operators.
+    relocate_ops: Enable relocation of ops.
+  '''
+  _os.environ['HB_OP_OPTIMIZATION_DISABLED'] = '0'
+  if logging_level is not None:
+    if 'TF_CPP_VMODULE' not in _os.environ:
+      _os.environ['TF_CPP_VMODULE'] = ''
+    if _os.environ['TF_CPP_VMODULE']:
+      _os.environ['TF_CPP_VMODULE'] += ','
+    _os.environ['TF_CPP_VMODULE'] += (
+      f'op_optimization={logging_level},'
+      f'replacing={logging_level},'
+      f'pruning={logging_level},'
+      f'relocation={logging_level},'
+      f'packing={logging_level},'
+      f'fusion={logging_level}')
+  if relocate_ops:
+    _os.environ['HB_OP_RELOCATION_ENABLED'] = '1'
diff --git a/hybridbackend/tensorflow/plugins/deeprec/__init__.py b/hybridbackend/tensorflow/plugins/deeprec/__init__.py
index 87d47d8f..11bc4c7b 100644
--- a/hybridbackend/tensorflow/plugins/deeprec/__init__.py
+++ b/hybridbackend/tensorflow/plugins/deeprec/__init__.py
@@ -23,24 +23,10 @@
 from hybridbackend.tensorflow.framework.context import Context as _ctx
 from hybridbackend.tensorflow.framework.options import DictOption as _dict
 
-_ = (
-  _ctx.get().options
-  .register('paiev_enabled', _dict(True))
-  .register('paiev_regularizer', _dict(None))
-  .register('paiev_caching_device', _dict(None))
-  .register('paiev_partitioner', _dict(None))
-  .register('paiev_validate_shape', _dict(True))
-  .register('paiev_custom_getter', _dict(None))
-  .register('paiev_constraint', _dict(None))
-  .register('paiev_steps_to_live', _dict(None))
-  .register('paiev_init_data_source', _dict(None))
-  .register('paiev_ev_option', _dict(None)))
-
 # pylint: disable=ungrouped-imports
 try:
-  from .ev import EmbeddingBackendDeepRecEV  # pylint: disable=unused-import
   from .ev import \
-    EmbeddingLookupRewritingForDeepRecEV  # pylint: disable=unused-import
+    EmbeddingWeightsRewritingForDeepRecEV  # pylint: disable=unused-import
 except ImportError:
   pass
 # pylint: enable=ungrouped-imports
diff --git a/hybridbackend/tensorflow/plugins/deeprec/ev.py b/hybridbackend/tensorflow/plugins/deeprec/ev.py
index c19cd2c1..ce77704f 100644
--- a/hybridbackend/tensorflow/plugins/deeprec/ev.py
+++ b/hybridbackend/tensorflow/plugins/deeprec/ev.py
@@ -20,204 +20,50 @@
 from __future__ import division
 from __future__ import print_function
 
-import six
-
-from tensorflow.python.ops import embedding_ops
-from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variable_scope as vs
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import tf_logging as logging
 
-from hybridbackend.tensorflow.feature_column.backend import EmbeddingBackend
-from hybridbackend.tensorflow.feature_column.default import \
-  EmbeddingBackendDefault
-from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.training.embedding import \
-  EmbeddingLookupRewriting
-
-
-class EmbeddingBackendDeepRecEV(EmbeddingBackendDefault):  # pylint: disable=useless-object-inheritance
-  r'''Embedding backend for EV in DeepRec.
-  '''
-  NAME = 'PAIEV'
-
-  def enabled(self, column):
-    r'''Enable EV for specific columns.
-    '''
-    enabled = (
-      Context.get().options.paiev_enabled[column])
-    if not enabled and self.num_buckets(column) is None:
-      logging.warning(
-        'Not possible to disable EV for column {column} with `None` buckets.')
-      return True
-    return enabled
-
-  def sharded(self, column):
-    r'''Whether the column should be sharded.
-    '''
-    if Context.get().world_size <= 1:
-      return False
-    if not self.enable_sharding:
-      return False
-    return True
-
-  def device(self, column):
-    r'''Device of the column weights.
-    '''
-    if not self.enabled(column):
-      return super().device(column)
-
-    emb_device = super().device(column)
-    if 'cpu' not in emb_device.lower():
-      logging.info('EV only supports CPU as embedding device')
-      return '/cpu:0'
-    return emb_device
-
-  def build(
-      self,
-      column,
-      name,
-      shape,
-      dtype=None,
-      trainable=True,
-      use_resource=True,
-      initializer=None,
-      collections=None,
-      layer=None):
-    r'''Creates the embedding lookup weight.
-    '''
-    if not self.enabled(column):
-      if len(shape) < 2:
-        raise ValueError(
-          'Non PAIEV column shall have a shape of at least 2 ranks')
-      return super().build(
-        column, name, shape,
-        dtype=dtype,
-        trainable=trainable,
-        use_resource=use_resource,
-        initializer=initializer,
-        collections=collections,
-        layer=layer)
+  EmbeddingWeightsRewriting
 
-    self._input_dtype = self.input_dtype(column)
-    options = Context.get().options
-    regularizer = options.paiev_regularizer[column]
-    caching_device = options.paiev_caching_device[column]
-    partitioner = options.paiev_partitioner[column]
-    validate_shape = options.paiev_validate_shape[column]
-    custom_getter = options.paiev_custom_getter[column]
-    constraint = options.paiev_constraint[column]
-    steps_to_live = options.paiev_steps_to_live[column]
-    init_data_source = options.paiev_init_data_source[column]
-    ev_option = options.paiev_ev_option[column]
-    if ev_option is None:
-      ev_option = variables.EmbeddingVariableOption()
 
-    var = vs.get_embedding_variable(
-      name,
-      shape[-1],
-      key_dtype=self._input_dtype,
-      value_dtype=dtype,
-      initializer=initializer,
-      regularizer=regularizer,
-      trainable=trainable,
-      collections=collections,
-      caching_device=caching_device,
-      partitioner=partitioner,
-      validate_shape=validate_shape,
-      custom_getter=custom_getter,
-      constraint=constraint,
-      steps_to_live=steps_to_live,
-      init_data_source=init_data_source,
-      ev_option=ev_option)
-    if hasattr(var, '_set_save_slice_info'):
-      var._set_save_slice_info(  # pylint: disable=protected-access
-        variables.Variable.SaveSliceInfo(
-          full_name=name,
-          full_shape=[Context.get().world_size, self.dimension(column)],
-          var_offset=[Context.get().rank, 0],
-          var_shape=shape))
-    elif isinstance(var, variables.PartitionedVariable):
-      for pvar in var:
-        pvar._set_save_slice_info(  # pylint: disable=protected-access
-          variables.Variable.SaveSliceInfo(
-            full_name=name,
-            full_shape=[Context.get().world_size, self.dimension(column)],
-            var_offset=[Context.get().rank, 0],
-            var_shape=pvar.shape))
-    else:
-      logging.warning(f'Restoring EV without elasticity: {var}')
-    return var
-
-  def lookup(self, column, weight, inputs, sharded=False, buffered=False):
-    r'''Lookup for embedding vectors.
-    '''
-    if not self.enabled(column):
-      return super().lookup(
-        column, weight, inputs,
-        sharded=sharded,
-        buffered=buffered)
-
-    inputs = math_ops.cast(inputs, self._input_dtype)
-    return embedding_ops.embedding_lookup(weight, inputs)
-
-  def update(self, column, weight, indexed_updates):
-    r'''Update embedding weight.
-    '''
-    if not self.enabled(column):
-      return super().update(column, weight, indexed_updates)
-
-    raise NotImplementedError
-
-  def weight_name(self, column):
-    r'''Name of the column weights.
-    '''
-    if self.sharded(column):
-      shard = Context.get().rank
-      return f'embedding_weights/part_{shard}'
-    return 'embedding_weights/part_0'
-
-  def weight_shared_name(self, column, var):
-    r'''Get shared name of the column weights from an variable.
-    '''
-    var_name = var.name.split(':')[0]
-    if self.sharded(column):
-      return var_name[:var_name.rfind('/part')]
-    return var_name
-
-
-EmbeddingBackend.register(EmbeddingBackendDeepRecEV())
-
-
-class EmbeddingLookupRewritingForDeepRecEV(EmbeddingLookupRewriting):  # pylint: disable=useless-object-inheritance
+class EmbeddingWeightsRewritingForDeepRecEV(EmbeddingWeightsRewriting):  # pylint: disable=useless-object-inheritance
   r'''Embedding lookup decorator for DeepRec EV.
   '''
   def __init__(self):
     super().__init__()
-    self._prev_lookup = None
     self._prev_get_embedding_variable = None
 
-  def build_unsharded_weights(self, fn, name, *args, **kwargs):
+  @property
+  def isdynamic(self):
+    r'''Whether embedding weights is dynamic.
+    '''
+    return True
+
+  def build_unsharded_weights(self, fn, *args, **kwargs):
     r'''Build unsharded embedding weights.
     '''
-    return fn(f'{name}/part_0', *args, **kwargs)
+    var_scope, var_store, name, *next_args = args
+    return fn(var_scope, var_store, f'{name}/part_0', *next_args, **kwargs)
 
-  def build_sharded_weights(self, shard, num_shards, fn, name, *args, **kwargs):
+  def build_sharded_weights(
+      self, shard, num_shards, shard_collections, fn, *args, **kwargs):
     r'''Build sharded embedding weights.
     '''
-    embedding_dim = kwargs.pop('embedding_dim', None)
+    kwargs['collections'] = shard_collections
+    embedding_dim = kwargs.get('shape', None)
     if embedding_dim is None:
-      try:
-        embedding_dim, = args
-      except ValueError as ex:
-        six.raise_from(
-          ValueError('missing embedding_dim for tf.get_embedding_variable'),
-          ex)
-    embedding_weights = fn(f'{name}/part_{shard}', embedding_dim, **kwargs)
+      raise ValueError('missing embedding_dim for tf.get_embedding_variable')
+    var_scope, var_store, name, *next_args = args
+    embedding_weights = fn(
+      var_scope, var_store, f'{name}/part_{shard}', *next_args, **kwargs)
+    full_name = embedding_weights.name.split(':')[0]
+    full_name = full_name[:full_name.rfind('/part')]
     if hasattr(embedding_weights, '_set_save_slice_info'):
       embedding_weights._set_save_slice_info(  # pylint: disable=protected-access
         variables.Variable.SaveSliceInfo(
-          full_name=name,
+          full_name=full_name,
           full_shape=[num_shards, embedding_dim],
           var_offset=[shard, 0],
           var_shape=embedding_weights.shape))
@@ -225,12 +71,13 @@ def build_sharded_weights(self, shard, num_shards, fn, name, *args, **kwargs):
       for pvar in embedding_weights:
         pvar._set_save_slice_info(  # pylint: disable=protected-access
           variables.Variable.SaveSliceInfo(
-            full_name=name,
+            full_name=full_name,
             full_shape=[num_shards, embedding_dim],
             var_offset=[shard, 0],
             var_shape=pvar.shape))
     else:
-      logging.warning(f'Embedding weights {name} cannot be saved correctly')
+      logging.warning(
+        f'Embedding weights {full_name} cannot be saved correctly')
 
     return embedding_weights
 
@@ -238,18 +85,10 @@ def begin(self):
     r'''Rewrites API.
     '''
     try:
-      import tensorflow as tf  # pylint: disable=import-outside-toplevel
-      self._prev_lookup = embedding_ops.embedding_lookup
-      embedding_ops.embedding_lookup = self.wraps_embedding_lookup(
-        embedding_ops.embedding_lookup)
-      tf.nn.embedding_lookup = self.wraps_embedding_lookup(
-        embedding_ops.embedding_lookup)
-
-      self._prev_get_embedding_variable = vs.get_embedding_variable
-      vs.get_embedding_variable = self.wraps_build_embedding_weights(
-        self._prev_get_embedding_variable)
-      tf.get_embedding_variable = self.wraps_build_embedding_weights(
-        self._prev_get_embedding_variable)
+      self._prev_get_embedding_variable = (
+        vs.VariableScope.get_embedding_variable)  # pylint: disable=protected-access
+      vs.VariableScope.get_embedding_variable = (  # pylint: disable=protected-access
+        self.wraps_build_embedding_weights(self._prev_get_embedding_variable))
     except:  # pylint: disable=bare-except
       pass
 
@@ -257,14 +96,10 @@ def end(self):
     r'''Revert API rewriting.
     '''
     try:
-      import tensorflow as tf  # pylint: disable=import-outside-toplevel
-      vs.get_embedding_variable = self._prev_get_embedding_variable
-      tf.get_embedding_variable = self._prev_get_embedding_variable
-
-      embedding_ops.embedding_lookup = self._prev_lookup
-      tf.nn.embedding_lookup = self._prev_lookup
+      vs.VariableScope.get_embedding_variable = (  # pylint: disable=protected-access
+        self._prev_get_embedding_variable)
     except:  # pylint: disable=bare-except
       pass
 
 
-EmbeddingLookupRewriting.register(EmbeddingLookupRewritingForDeepRecEV)
+EmbeddingWeightsRewriting.register(EmbeddingWeightsRewritingForDeepRecEV)
diff --git a/hybridbackend/tensorflow/training/__init__.py b/hybridbackend/tensorflow/training/__init__.py
index 7024f45a..e38b79e1 100644
--- a/hybridbackend/tensorflow/training/__init__.py
+++ b/hybridbackend/tensorflow/training/__init__.py
@@ -45,16 +45,14 @@
 from hybridbackend.tensorflow.training.sync_replicas_optimizer import \
   SyncReplicasOptimizer
 from hybridbackend.tensorflow.training.variables import \
-  EmbeddingLookupRewritingForVariables as _patch_variables
+  EmbeddingWeightsRewritingForVariables as _patch_variables
 
 _ = (
   _ctx.get().options
   .register('grad_lazy_sync', False, env='HB_GRAD_LAZY_SYNC')
-  .register('grad_nbuckets', 1, env='HB_GRAD_NBUCKETS')
   .register('sharding', False)
   .register('batch_size', -1)
   .register('model_dir', None)
-  .register('eval_dir', None)
   .register('keep_checkpoint_max', None)
   .register('keep_checkpoint_every_n_hours', None)
   .register('mode', _mode_keys.TRAIN))
diff --git a/hybridbackend/tensorflow/training/config.py b/hybridbackend/tensorflow/training/config.py
deleted file mode 100644
index 13faa54c..00000000
--- a/hybridbackend/tensorflow/training/config.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''ConfigProto related functions.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.core.protobuf import config_pb2
-from tensorflow.python.framework import device as pydev
-
-from hybridbackend.tensorflow.framework.context import Context
-
-
-def configure(*args, prototype=None, **kwargs):
-  r'''Creates ConfigProto.
-  '''
-  ctx = Context.get()
-  if not prototype:
-    kwargs.setdefault('allow_soft_placement', True)
-    prototype = config_pb2.ConfigProto(*args, **kwargs)
-  prototype.graph_options.optimizer_options.global_jit_level = (
-    config_pb2.OptimizerOptions.OFF)
-  prototype.gpu_options.allow_growth = True
-  prototype.gpu_options.force_gpu_compatible = True
-  chief = pydev.DeviceSpec.from_string(ctx.devices[0])
-  del prototype.device_filters[:]
-  prototype.device_filters.append(f'/job:{chief.job}/task:{chief.task}')
-  prototype.device_filters.append(
-    f'/job:{ctx.task_type}/task:{ctx.task_id}')
-  prototype.experimental.collective_group_leader = (
-    f'/job:{chief.job}/replica:{chief.replica}/task:{chief.task}')
-  return prototype
diff --git a/hybridbackend/tensorflow/training/embedding.py b/hybridbackend/tensorflow/training/embedding.py
index cdc1a705..22fadce6 100644
--- a/hybridbackend/tensorflow/training/embedding.py
+++ b/hybridbackend/tensorflow/training/embedding.py
@@ -22,19 +22,24 @@
 
 import abc
 
+from tensorflow.python.feature_column.feature_column import \
+  _SharedEmbeddingColumn as _shared_emb
 from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import embedding_ops
 
-from hybridbackend.tensorflow.distribute.embedding import \
-  ParallelEmbeddingLookupContext
+from hybridbackend.tensorflow.distribute.collective import Collective
+from hybridbackend.tensorflow.distribute.partition.ops import \
+  partition_by_modulo
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.ops import GraphKeys
 from hybridbackend.tensorflow.framework.ops import ModeKeys
 from hybridbackend.tensorflow.framework.rewriting import GraphRewriting
 
 
-class EmbeddingLookupRewriting(GraphRewriting):  # pylint: disable=useless-object-inheritance
-  r'''Rewriting embedding lookups.
+class EmbeddingWeightsRewriting(GraphRewriting):  # pylint: disable=useless-object-inheritance
+  r'''Rewriting embedding weights.
   '''
   @classmethod
   def register(cls, rewriting):
@@ -45,6 +50,12 @@ def register(cls, rewriting):
     '''
     return GraphRewriting.register(rewriting)
 
+  @property
+  def isdynamic(self):
+    r'''Whether embedding weights is dynamic.
+    '''
+    return False
+
   @property
   def should_shard(self):
     r'''Whether embedding weights sharding should be enabled.
@@ -55,14 +66,6 @@ def should_shard(self):
       and ctx.options.sharding
       and ctx.options.mode != ModeKeys.PREDICT)
 
-  @property
-  def num_shards(self):
-    return Context.get().world_size
-
-  @property
-  def shard(self):
-    return Context.get().rank
-
   def _nosharding(self, fn):
     r'''Wraps fn without sharding.
     '''
@@ -76,93 +79,169 @@ def _sharded_embedding_weights_collections(self, collections):
     '''
     minimal_collections = [
       ops.GraphKeys.GLOBAL_VARIABLES,
-      GraphKeys.SHARDED_VARIABLES,
-      self.__class__.__name__]
+      GraphKeys.SHARDED_VARIABLES]
     if collections is None:
-      collections = list(minimal_collections)
-    else:
-      collections.extend(minimal_collections)
-      collections = list(set(collections))
-    return collections
+      return list(minimal_collections)
+    collections_copy = list(collections)
+    collections_copy.extend(minimal_collections)
+    if self.isdynamic:
+      collections_copy.append(GraphKeys.DYNAMIC_VARIABLES)
+    return list(set(collections_copy))
+
+  def wraps_build_embedding_weights(self, fn):
+    r'''Rewrites weights building.
+    '''
+    def _wrapped_build_weights(*args, **kwargs):
+      r'''Build embedding weights.
+      '''
+      if not self.should_shard:
+        return self.build_unsharded_weights(fn, *args, **kwargs)
+
+      shard_collections = kwargs.get('collections', None)
+      shard_collections = self._sharded_embedding_weights_collections(
+        shard_collections)
+      return self.build_sharded_weights(
+        Context.get().rank, Context.get().world_size,
+        shard_collections, self._nosharding(fn),
+        *args, **kwargs)
+    return _wrapped_build_weights
+
+  def build_unsharded_weights(self, fn, *args, **kwargs):
+    r'''Build unsharded embedding weights.
+    '''
+    return fn(*args, **kwargs)
+
+  def build_sharded_weights(
+      self, shard, num_shards, shard_collections, fn, *args, **kwargs):
+    r'''Build sharded embedding weights.
+    '''
+    del shard
+    del num_shards
+    del shard_collections
+    return fn(*args, **kwargs)
+
+  @abc.abstractmethod
+  def begin(self):
+    r'''Rewrites API.
+    '''
 
-  def _is_sharded(self, weights):
+  @abc.abstractmethod
+  def end(self):
+    r'''Revert API rewriting.
+    '''
+
+
+class EmbeddingLookupRewriting(GraphRewriting):
+  r'''Embedding lookup rewriting for tf.nn.embedding_lookup.
+  '''
+  def __init__(self):
+    super().__init__()
+    self._prev_lookup = None
+    self._prev_shared_emb_get_dense_tensor_internal = None
+    self._local_world_size = Context.get().local_world_size
+    self._num_nodes = Context.get().world_size // self._local_world_size
+
+  def issharded(self, weights):
     r'''Check whether the embedding weights are sharded.
     '''
-    sharded_weights_set = ops.get_default_graph().get_collection_ref(
-      self.__class__.__name__)
-    if weights in sharded_weights_set:
+    sharded_variables = ops.get_default_graph().get_collection_ref(
+      GraphKeys.SHARDED_VARIABLES)
+    if weights in sharded_variables:
       return True
     if isinstance(weights, (list, tuple)) and len(weights) == 1:
       weights = weights[0]
       if isinstance(weights, ops.Tensor):
         vname = weights.name.split('/read')[0]
-        for v in sharded_weights_set:
+        for v in sharded_variables:
           if vname == v.name.split(':')[0]:
             return True
     return False
 
-  def wraps_build_embedding_weights(self, fn):
-    r'''Rewrites weights building.
-    '''
-    def _wrapped_build_weights(name, *args, **kwargs):
-      r'''Build embedding weights.
-      '''
-      if not self.should_shard:
-        return self.build_unsharded_weights(fn, name, *args, **kwargs)
-
-      collections = kwargs.pop('collections', None)
-      kwargs['collections'] = self._sharded_embedding_weights_collections(
-        collections)
-      return self.build_sharded_weights(
-        self.shard, self.num_shards, self._nosharding(fn),
-        name, *args, **kwargs)
-    return _wrapped_build_weights
-
   def wraps_embedding_lookup(self, fn):
     r'''Rewrites embedding_lookup.
     '''
     def _wrapped_embedding_lookup(params, ids, **kwargs):
       r'''Looks up `ids` in a list of embedding tensors.
       '''
-      if not self._is_sharded(params):
+      if not self.issharded(params):
         with Context.scope(sharding=False):
           return fn(params, ids, **kwargs)
 
       current_device = control_flow_ops.no_op().device
+      num_shards = Context.get().world_size
       with Context.scope(sharding=False):
         with ops.device(Context.get().devices[Context.get().rank]):
-          ctx = ParallelEmbeddingLookupContext(current_device)
-          ids = ctx.exchange_ids(ids)
-          ids = self.build_sharded_ids(ids)
-          with ops.device(ctx.device):
-            embeddings = fn(params, ids, **kwargs)
-        return ctx.exchange_embeddings(embeddings)
+          ids_shards, ids_sizes, shard_index = partition_by_modulo(
+            ids, num_shards, name='shard_partition')
+          shard_ids, shard_sizes = Collective.get().alltoall(
+            ids_shards, sizes=ids_sizes)
+          shard_ids, shard_unique_index = array_ops.unique(
+            shard_ids, name='shard_unique')
+          if params not in ops.get_collection_ref(GraphKeys.DYNAMIC_VARIABLES):
+            shard_ids = shard_ids // num_shards
+          with ops.device(current_device):
+            embeddings = fn(params, shard_ids, **kwargs)
+            dimension = int(embeddings.shape[-1])
+          embeddings = array_ops.gather(
+            embeddings, shard_unique_index,
+            name='shard_unique_restore')
+          embeddings, _ = Collective.get().alltoall(
+            embeddings,
+            sizes=shard_sizes,
+            common_shape=[dimension])
+          embeddings = array_ops.gather(
+            embeddings, shard_index,
+            name='shard_stitch')
+        return embeddings
 
     return _wrapped_embedding_lookup
 
-  def build_sharded_ids(self, ids):
-    r'''Shard IDs to lookup sharded embedding weights.
+  def wraps_shared_embedding_get_dense_tensor_internal(self, fn):
+    r'''Rewrites to prevent an incorrect embedding_shape.
     '''
-    return ids
-
-  def build_unsharded_weights(self, fn, name, *args, **kwargs):
-    r'''Build unsharded embedding weights.
-    '''
-    return fn(name, *args, **kwargs)
-
-  def build_sharded_weights(self, shard, num_shards, fn, name, *args, **kwargs):
-    r'''Build sharded embedding weights.
-    '''
-    del shard
-    del num_shards
-    return fn(name, *args, **kwargs)
+    def _wrapped_shared_embedding_get_dense_tensor_internal(
+        cls, *args, **kwargs):
+      r'''Get dense tensors in shared embedding column.
+      '''
+      shared_embedding_collection = ops.get_collection(
+        cls.shared_embedding_collection_name)
+      if shared_embedding_collection:
+        embedding_shape = (cls.categorical_column._num_buckets, cls.dimension)  # pylint: disable=protected-access
+        embedding_weights = shared_embedding_collection[0]
+        prev_embedding_shape = embedding_weights.get_shape
+        embedding_weights.get_shape = lambda: embedding_shape
+        ret = fn(cls, *args, **kwargs)
+        embedding_weights.get_shape = prev_embedding_shape
+        return ret
+      return fn(cls, *args, **kwargs)
+    return _wrapped_shared_embedding_get_dense_tensor_internal
 
-  @abc.abstractmethod
   def begin(self):
     r'''Rewrites API.
     '''
+    import tensorflow as tf  # pylint: disable=import-outside-toplevel
+    self._prev_lookup = embedding_ops.embedding_lookup
+
+    embedding_ops.embedding_lookup = self.wraps_embedding_lookup(
+      embedding_ops.embedding_lookup)
+    tf.nn.embedding_lookup = self.wraps_embedding_lookup(
+      embedding_ops.embedding_lookup)
+
+    # pylint: disable=protected-access
+    self._prev_shared_emb_get_dense_tensor_internal = (
+      _shared_emb._get_dense_tensor_internal)
+    _shared_emb._get_dense_tensor_internal = (
+      self.wraps_shared_embedding_get_dense_tensor_internal(
+        self._prev_shared_emb_get_dense_tensor_internal))
 
-  @abc.abstractmethod
   def end(self):
     r'''Revert API rewriting.
     '''
+    import tensorflow as tf  # pylint: disable=import-outside-toplevel
+    embedding_ops.embedding_lookup = self._prev_lookup
+    tf.nn.embedding_lookup = self._prev_lookup
+    _shared_emb._get_dense_tensor_internal = (  # pylint: disable=protected-access
+      self._prev_shared_emb_get_dense_tensor_internal)
+
+
+GraphRewriting.register(EmbeddingLookupRewriting)
diff --git a/hybridbackend/tensorflow/training/evaluation.py b/hybridbackend/tensorflow/training/evaluation.py
index ef2857f0..37e919d3 100644
--- a/hybridbackend/tensorflow/training/evaluation.py
+++ b/hybridbackend/tensorflow/training/evaluation.py
@@ -115,6 +115,7 @@ def __init__(
       self, fn,
       steps=100,
       every_n_iter=1000,
+      summary_dir=None,
       history=None):
     r'''Evaluates specific function.
 
@@ -122,6 +123,7 @@ def __init__(
       fn: Function returns update_op, metric ops and hooks.
       steps: Number of steps for which to evaluate model.
       every_n_iter: `int`, runs the evaluator once every N training iteration.
+      summary_dir: Directory for summaries.
       history: History of eval metrics. history should support `append` method.
     Raises:
       ValueError: if `every_n_iter` is non-positive or it's not a single machine
@@ -132,6 +134,7 @@ def __init__(
     self._fn = fn
     self._steps = steps
     self._every_n_iter = every_n_iter
+    self._summary_dir = summary_dir
     self._history = history
 
   def begin(self):
@@ -148,7 +151,6 @@ def begin(self):
     self._timer = basic_session_run_hooks.SecondOrStepTimer(
       every_steps=self._every_n_iter)
     self._timer.reset()
-    self._summary_dir = Context.get().options.eval_dir or '.'
 
     with ops.device(device_function), reuse_variables(vs.AUTO_REUSE):
       with scope(mode=ModeKeys.EVAL, comm_pool_name=ModeKeys.EVAL):
@@ -251,7 +253,7 @@ def _run(self, run_context, fetches):
     return eval_results
 
   def _write_dict_to_summary(self, dictionary):
-    r'''Write evaluation results to eval_dir.
+    r'''Write evaluation results to summary directory.
     '''
     current_global_step = dictionary[ops.GraphKeys.GLOBAL_STEP]
     prev_np_printoptions = np.get_printoptions()
diff --git a/hybridbackend/tensorflow/distribute/gradient.py b/hybridbackend/tensorflow/training/gradient.py
similarity index 50%
rename from hybridbackend/tensorflow/distribute/gradient.py
rename to hybridbackend/tensorflow/training/gradient.py
index ab80c4ee..5d92fdc9 100644
--- a/hybridbackend/tensorflow/distribute/gradient.py
+++ b/hybridbackend/tensorflow/training/gradient.py
@@ -21,12 +21,10 @@
 from __future__ import print_function
 
 from tensorflow.python.framework import ops
-from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import tf_logging as logging
 
-from hybridbackend.tensorflow.distribute.communicator_pool import \
-  CommunicatorPool
+from hybridbackend.tensorflow.distribute.collective import Collective
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.ops import GraphKeys
 
@@ -40,18 +38,24 @@ def _tensor_size(tensor):
   return -1
 
 
-def _log(tensors, sparse_tensors, shard_tensors):
+def _log(
+    dense_vars, dense_grads,
+    sparse_vars, sparse_grads,
+    sharded_vars, sharded_grads):
   r'''Log values of gradients.
   Args:
-    tensors: list of dense gradients.
-    sparse_tensors: list of sparse gradients.
-    shard_tensors: list of tensors skipped in aggregation.
+    dense_vars: list of dense variables.
+    dense_grads: list of dense gradients.
+    sparse_vars: list of sparse variables.
+    sparse_grads: list of sparse gradients.
+    sharded_vars: list of sharded variables.
+    sharded_grads: list of sharded gradients.
   '''
-  sizes = [_tensor_size(t) for t in tensors]
+  sizes = [_tensor_size(t) for t in dense_grads]
   nonempty_sizes = [s for s in sizes if s > 0]
   sparse_sizes = [
     _tensor_size(s.values) + _tensor_size(s.indices)
-    for s in sparse_tensors]
+    for s in sparse_grads]
   nonempty_sparse_sizes = [s for s in sparse_sizes if s > 0]
   grad_mbs = sum(nonempty_sizes) / 1024.0 / 1024.0
   sparse_grad_mbs = sum(nonempty_sparse_sizes) / 1024.0 / 1024.0
@@ -60,99 +64,14 @@ def _log(tensors, sparse_tensors, shard_tensors):
     f'({">" if len(sizes) > len(nonempty_sizes) else ""}{grad_mbs:.2f}MB) '
     f'and {len(sparse_sizes)} sparse gradients '
     f'({">" if len(sparse_sizes) > len(nonempty_sparse_sizes) else ""}'
-    f'{sparse_grad_mbs:.2f}MB), '
-    f'skip {len(shard_tensors)} aggregated gradients')
-
-
-def _flatten(tensors):
-  r'''Flatten list of tensors into one buffer tensor.
-  Args:
-    tensors: tensors to flatten.
-  Returns:
-    flattened: flattened buffer tensor.
-    shapes: shapes of tensors.
-    sizes: sizes of tensors.
-  '''
-  ftensors = [array_ops.reshape(t, [-1]) for t in tensors]
-  shapes = [array_ops.shape(t) for t in tensors]
-  sizes = [array_ops.reshape(array_ops.shape(t), []) for t in ftensors]
-  return array_ops.concat(ftensors, 0), shapes, sizes
-
-
-def _deflatten(flattened, shapes, sizes):
-  r'''Get original tensor list from flattened buffer tensor.
-  Args:
-    flattened: flattened buffer tensor.
-    shapes: shapes of tensors.
-    sizes: sizes of tensors.
-  Returns:
-    original tensor list.
-  '''
-  ftensors = array_ops.split(flattened, sizes)
-  return [array_ops.reshape(t, shapes[i]) for i, t in enumerate(ftensors)]
-
-
-def _bucketize(grads, num_buckets):
-  r'''Split gradients into buckets.
-  Args:
-    grads: gradients to split.
-  Returns:
-    list of tensor buckets.
-  '''
-  if not grads:
-    return []
-  if len(grads) == 1:
-    return [grads]
-  # Pass 1: split by dtypes.
-  dtype_buckets = []
-  prev_dtype = None
-  for g in grads:
-    if g.dtype != prev_dtype:
-      dtype_buckets.append([g])
-      prev_dtype = g.dtype
-    else:
-      dtype_buckets[-1].append(g)
-  num_dtypes = len(dtype_buckets)
-  if num_dtypes >= num_buckets:
-    if num_dtypes > num_buckets:
-      logging.warning(f'num_buckets is increased to {num_dtypes}')
-    return dtype_buckets
-  # Pass 2: split by sizes.
-  # Pass 2.1: Calculate number of splits for different dtypes.
-  dtype_bucket_sizes = [len(b) for b in dtype_buckets]
-  dtype_bucket_total_size = sum(dtype_bucket_sizes)
-  bucket_num_splits = [
-    int(num_buckets * cnt / dtype_bucket_total_size)
-    for cnt in dtype_bucket_sizes]
-  bucket_num_splits = [n if n > 0 else 1 for n in bucket_num_splits]
-  bucket_num_splits[0] += (num_buckets - sum(bucket_num_splits))
-  buckets = []
-  for bucket_id, bucket in enumerate(dtype_buckets):
-    # Pass 2.2: For each dtype, calculate split size.
-    bucket_info = [(g.dtype.size, g.shape.num_elements()) for g in bucket]
-    bucket_sizes_or_none = [s * n if n else None for s, n in bucket_info]
-    bucket_sizes_without_none = [s for s in bucket_sizes_or_none if s]
-    bucket_min_total_size = sum(bucket_sizes_without_none)
-    bucket_mean_size = bucket_min_total_size / len(bucket_sizes_without_none)
-    bucket_sizes = [
-      s if s else bucket_mean_size for s in bucket_sizes_or_none]
-    if bucket_num_splits[bucket_id] == 1:
-      split_size = sum(bucket_sizes)
-    else:
-      split_size = sum(bucket_sizes) / (bucket_num_splits[bucket_id] - 1)
-    # Pass 2.3: For each dtype, split buckets by split size.
-    ext_buckets = []
-    prev_size = 0
-    for gid, g in enumerate(bucket):
-      if prev_size == 0 or prev_size + bucket_sizes[gid] > split_size:
-        ext_buckets.append([g])
-        prev_size = bucket_sizes[gid]
-      else:
-        ext_buckets[-1].append(g)
-        prev_size += bucket_sizes[gid]
-    buckets.extend(ext_buckets)
-  buckets = [g for g in buckets if g]
-  return buckets
+    f'{sparse_grad_mbs:.2f}MB) in optimizer, '
+    f'skip {len(sharded_grads)} gradients')
+  for v in dense_vars:
+    logging.info(f'Replicated weights: {v.name.split(":")[0]}')
+  for v in sparse_vars:
+    logging.info(f'Replicated weights: *{v.name.split(":")[0]}')
+  for v in sharded_vars:
+    logging.info(f'Sharded weights: {v.name.split(":")[0]}')
 
 
 def _mean(grads, devices):
@@ -197,25 +116,34 @@ def _pack_grads_and_vars(
   return grads_and_vars
 
 
-def _aggregate(replica_grads, shard_grads, num_buckets):
+def _aggregate(replicated_vars, replicated_grads, sharded_vars, sharded_grads):
   r'''Aggregate gradients.
   '''
-  if not replica_grads:
+  if not replicated_grads:
     return []
 
+  dense_vars = []
   dense_grads = []
   dense_grad_indices = []
+  sparse_vars = []
   sparse_grads = []
   sparse_grad_indices = []
-  for idx, g in enumerate(replica_grads):
-    if isinstance(g, ops.IndexedSlices):
+  for idx, g in enumerate(replicated_grads):
+    small_weights = ops.get_collection_ref(GraphKeys.TRAINABLE_REPLICATED_SMALL)
+    if (isinstance(g, ops.IndexedSlices)
+        and replicated_vars[idx] not in small_weights):
+      sparse_vars.append(replicated_vars[idx])
       sparse_grads.append(g)
       sparse_grad_indices.append(idx)
     else:
+      dense_vars.append(replicated_vars[idx])
       dense_grads.append(ops.convert_to_tensor(g))
       dense_grad_indices.append(idx)
-  _log(dense_grads, sparse_grads, shard_grads)
-  dense_aggregated = _aggregate_dense(dense_grads, num_buckets)
+  _log(
+    dense_vars, dense_grads,
+    sparse_vars, sparse_grads,
+    sharded_vars, sharded_grads)
+  dense_aggregated = _aggregate_dense(dense_grads)
   sparse_aggregated = _aggregate_sparse(sparse_grads)
   indexed_aggregated = []
   for i, g in enumerate(dense_aggregated):
@@ -226,38 +154,10 @@ def _aggregate(replica_grads, shard_grads, num_buckets):
   return aggregated
 
 
-def _aggregate_dense(grads, num_buckets):
+def _aggregate_dense(grads):
   r'''Aggregate dense gradients.
   '''
-  if not grads:
-    return []
-  if len(grads) == 1:
-    aggregated = CommunicatorPool.get().allreduce(grads[0], trainable=False)
-    return [aggregated]
-  # Bucketize gradients.
-  buckets = _bucketize(grads, num_buckets)
-  # Flatten buckets.
-  bucket_tensor_info = []
-  for bucket_id, bucket in enumerate(buckets):
-    with ops.name_scope(f'buckets/{bucket_id}'):
-      bucket_tensor_info.append(_flatten(bucket))
-  bucket_tensors, bucket_tensor_shapes, bucket_tensor_sizes = zip(
-    *bucket_tensor_info)
-  # Aggregate on flattened tensors.
-  bucket_aggregations = []
-  for bucket_id, flattened in enumerate(bucket_tensors):
-    aggregated = CommunicatorPool.get().allreduce(flattened, trainable=False)
-    bucket_aggregations.append(aggregated)
-  # Unflatten buckets.
-  bucket_grads = [
-    _deflatten(
-      g,
-      bucket_tensor_shapes[bucket_id],
-      bucket_tensor_sizes[bucket_id])
-    for bucket_id, g in enumerate(bucket_aggregations)]
-  # Unbucketize gradients.
-  aggregated_grads = [g for bucket in bucket_grads for g in bucket]
-  return aggregated_grads
+  return [Collective.get().allreduce(g) for g in grads]
 
 
 def _aggregate_sparse(grads):
@@ -267,10 +167,8 @@ def _aggregate_sparse(grads):
     return []
   aggregated_grads = []
   for g in grads:
-    aggregated_values = CommunicatorPool.get().allgatherv(
-      g.values, trainable=False)
-    aggregated_indices = CommunicatorPool.get().allgatherv(
-      g.indices, trainable=False)
+    aggregated_values = Collective.get().allgather(g.values)
+    aggregated_indices = Collective.get().allgather(g.indices)
     aggregated_grads.append(
       ops.IndexedSlices(
         aggregated_values,
@@ -279,7 +177,7 @@ def _aggregate_sparse(grads):
   return aggregated_grads
 
 
-def aggregate_gradients(grads_and_vars, num_buckets=1):
+def aggregate_gradients(grads_and_vars):
   r'''Aggregate gradients and variables.
   '''
   devices = Context.get().devices
@@ -307,12 +205,14 @@ def aggregate_gradients(grads_and_vars, num_buckets=1):
       replica_grad_indices.append(i)
       replica_grads.append(gv[0])
       replica_vars_to_optimize.append(gv[1])
-  replica_grads = _aggregate(replica_grads, shard_grads, num_buckets)
+  replica_grads = _aggregate(
+    replica_vars_to_optimize, replica_grads,
+    sharded_vars_to_optimize, shard_grads)
   ops.get_collection_ref(GraphKeys.TRAINABLE_REPLICATED).extend(
     replica_vars_to_optimize)
   ops.get_collection_ref(GraphKeys.TRAINABLE_SHARDED).extend(
     sharded_vars_to_optimize)
   return _pack_grads_and_vars(
-    _mean(replica_grads, devices), _mean(shard_grads, devices),
+    _mean(replica_grads, devices), shard_grads,
     replica_vars_to_optimize, sharded_vars_to_optimize,
     replica_grad_indices, shard_grad_indices)
diff --git a/hybridbackend/tensorflow/training/optimizer.py b/hybridbackend/tensorflow/training/optimizer.py
index 913fcfb8..45138473 100644
--- a/hybridbackend/tensorflow/training/optimizer.py
+++ b/hybridbackend/tensorflow/training/optimizer.py
@@ -32,13 +32,13 @@
 from tensorflow.python.training import distribution_strategy_context
 from tensorflow.python.training import training
 
-from hybridbackend.tensorflow.distribute.gradient import aggregate_gradients
-from hybridbackend.tensorflow.distribute.pubsub import PubSub
+from hybridbackend.tensorflow.distribute.rpc import RpcCollective
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.ops import GraphKeys
 from hybridbackend.tensorflow.framework.ops import ModeKeys
 from hybridbackend.tensorflow.framework.rewriting import GraphRewriting
 from hybridbackend.tensorflow.framework.rewriting import SessionRunRewriting
+from hybridbackend.tensorflow.training.gradient import aggregate_gradients
 
 
 class HybridBackendOptimizerBase(object):  # pylint: disable=useless-object-inheritance
@@ -48,22 +48,18 @@ class HybridBackendOptimizerBase(object):  # pylint: disable=useless-object-inhe
 
 def wraps_optimizer(
     cls,
-    num_buckets=None,
     aggregation=None):
   r'''Decorator to create hybridbackend optimizer class.
 
   Args:
     optimizer_type: The actual optimizer type that will be used to compute and
       apply the gradients. Must be one of the Optimizer classes.
-    num_buckets: Max number of gradient groups.
     aggregation: Aggregate gradients inside `compute_gradients` or
       `apply_gradients`.
 
   Returns:
     hb_optimizer_type: The hybridbackend optimizer type for `optimizer_type`
   '''
-  if num_buckets is None:
-    num_buckets = Context.get().options.grad_nbuckets
   if aggregation is None:
     aggregation = (
       'apply_gradients'
@@ -124,7 +120,7 @@ def _aggregate_gradients(self, grads_and_vars):
         # Do nothing if this function resides in a distribution context.
         return grads_and_vars
 
-      return aggregate_gradients(grads_and_vars, num_buckets)
+      return aggregate_gradients(grads_and_vars)
 
     def compute_gradients(self, *args, **kwargs):
       r'''Compute gradients of "loss" for the variables in "var_list".
@@ -222,7 +218,7 @@ def begin(self):
           ops.get_collection_ref(GraphKeys.TRAINABLE_REPLICATED)))
       for v in replicated_variables:
         try:
-          self._set_initializer(v, Context.get().devices)
+          self._set_initializer(v)
         except:  # pylint: disable=bare-except
           pass
 
@@ -250,12 +246,13 @@ def _get_initializer_op(self, var):
     return state_ops.assign(var._variable, self._get_initial_value(var)).op
     # pylint:enable=protected-access
 
-  def _set_initializer(self, var, devices, name=None):
+  def _set_initializer(self, var, name=None):
     r'''Initialize variables.
     '''
     if name is None:
       name = var.name.split(':')[0]
-    rank = Context.get().current_index()
+    devices = Context.get().devices
+    rank = Context.get().rank
     with ops.name_scope('initializers/'):
       with ops.name_scope(f'{name}/initializer'):
         with ops.control_dependencies(None):
@@ -268,11 +265,12 @@ def _set_initializer(self, var, devices, name=None):
             # pylint:disable=protected-access
             var._initial_value = initial_value
             if len(devices) > 1:
-              var._initial_value = PubSub(devices, rank=rank)(
-                lambda: initial_value,
-                initial_value.shape,
-                initial_value.dtype,
-                name=f'{name}_pubsub')
+              var._initial_value = (
+                RpcCollective(devices, rank).broadcast(
+                  lambda: initial_value,
+                  initial_value.dtype,
+                  initial_value.shape,
+                  name=f'{name}_initial_value/rpc_broadcast'))
             var._initializer_op = self._get_initializer_op(var)
             # pylint:enable=protected-access
 
diff --git a/hybridbackend/tensorflow/training/saver.py b/hybridbackend/tensorflow/training/saver.py
index 586eb7c0..48ead17a 100644
--- a/hybridbackend/tensorflow/training/saver.py
+++ b/hybridbackend/tensorflow/training/saver.py
@@ -30,6 +30,7 @@
 from tensorflow.python.ops import gen_io_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.training import checkpoint_management
 from tensorflow.python.training import checkpoint_utils
 from tensorflow.python.training import saver
@@ -134,6 +135,10 @@ def _AddShardedSaveOpsForV2(self, checkpoint_prefix, per_device):
           saveables = [
             s for s in saveables
             if s.op in sharded_saveables or s in sharded_saveables]
+        logging.vlog(
+          1,
+          f'Saving {len(saveables)} saveables for shard {shard} at process '
+          f'{self._rank}: {[s.name for s in saveables]}')
         last_device = device
         with ops.device(device):
           with ops.device('/cpu:0'):
diff --git a/hybridbackend/tensorflow/training/server.py b/hybridbackend/tensorflow/training/server.py
index 9c7d282f..8cb6b47d 100644
--- a/hybridbackend/tensorflow/training/server.py
+++ b/hybridbackend/tensorflow/training/server.py
@@ -23,9 +23,9 @@
 from tensorflow.python.training import monitored_session as _monitored_session
 from tensorflow.python.training import server_lib
 
+from hybridbackend.tensorflow.framework.config import wraps_session_config
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.rewriting import scope
-from hybridbackend.tensorflow.training.config import configure
 
 
 class HybridBackendServerBase(object):  # pylint: disable=useless-object-inheritance
@@ -47,21 +47,21 @@ class HybridBackendServer(cls, HybridBackendServerBase):
     @classmethod
     def get(class_):
       if class_._default is None:
-        class_._default = class_()
+        class_._default = class_(None)
       return class_._default
 
-    def __init__(self, server_or_cluster_def=None, **kwargs):
+    def __init__(self, server_or_cluster_def, **kwargs):
       r'''Creates a new server with the given definition.
       '''
       if server_or_cluster_def is None:
         server_or_cluster_def = Context.get().cluster_spec
+        kwargs['job_name'] = Context.get().task_type
+        kwargs['task_index'] = Context.get().task_id
       if server_or_cluster_def is None:
         self._is_local = True
         return
       self._is_local = False
-      kwargs['job_name'] = Context.get().task_type
-      kwargs['task_index'] = Context.get().task_id
-      kwargs['config'] = configure(prototype=kwargs.pop('config', None))
+      kwargs['config'] = wraps_session_config(kwargs.pop('config', None))
       super().__init__(server_or_cluster_def, **kwargs)
 
     @property
diff --git a/hybridbackend/tensorflow/training/session.py b/hybridbackend/tensorflow/training/session.py
index 86eee10f..6ca51f0b 100644
--- a/hybridbackend/tensorflow/training/session.py
+++ b/hybridbackend/tensorflow/training/session.py
@@ -26,11 +26,11 @@
 from tensorflow.python.training import monitored_session as _monitored_session
 from tensorflow.python.training import training
 
+from hybridbackend.tensorflow.framework.config import wraps_session_config
 from hybridbackend.tensorflow.framework.context import Context
 from hybridbackend.tensorflow.framework.device import device_function
 from hybridbackend.tensorflow.framework.rewriting import GraphRewriting
 from hybridbackend.tensorflow.framework.rewriting import SessionRunRewriting
-from hybridbackend.tensorflow.training.config import configure
 from hybridbackend.tensorflow.training.hooks import Policy
 from hybridbackend.tensorflow.training.server import Server
 
@@ -105,7 +105,7 @@ def HybridBackendMonitoredTrainingSession(*args, **kwargs):  # pylint: disable=i
           output_dir=summary_dir))
     kwargs['hooks'] = hooks
     kwargs['chief_only_hooks'] = chief_only_hooks
-    kwargs['config'] = configure(prototype=kwargs.pop('config', None))
+    kwargs['config'] = wraps_session_config(kwargs.pop('config', None))
     kwargs['is_chief'] = True
     args = list(args)
     if args:
@@ -119,9 +119,7 @@ def HybridBackendMonitoredTrainingSession(*args, **kwargs):  # pylint: disable=i
         master = Server.get().target
       kwargs['master'] = master
 
-    with ops.device(device_function), Context.scope(
-        model_dir=checkpoint_dir,
-        eval_dir=Context.get().options.eval_dir or summary_dir):
+    with ops.device(device_function), Context.scope(model_dir=checkpoint_dir):
       prev_monitored_session = _monitored_session.MonitoredSession
       _monitored_session.MonitoredSession = wraps_monitored_session(
         prev_monitored_session,
diff --git a/hybridbackend/tensorflow/training/variables.py b/hybridbackend/tensorflow/training/variables.py
index d7d44a03..3398f497 100644
--- a/hybridbackend/tensorflow/training/variables.py
+++ b/hybridbackend/tensorflow/training/variables.py
@@ -27,16 +27,17 @@
 from tensorflow.python.keras.backend import reset_uids as reset_keras_uids
 from tensorflow.python.keras.engine import base_layer
 from tensorflow.python.layers import base
-from tensorflow.python.ops import embedding_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.ops import variable_scope as vs
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import tf_logging as logging
 
+from hybridbackend.tensorflow.framework.context import Context
+from hybridbackend.tensorflow.framework.ops import GraphKeys
 from hybridbackend.tensorflow.framework.rewriting import GraphRewriting
 from hybridbackend.tensorflow.training.embedding import \
-  EmbeddingLookupRewriting
+  EmbeddingWeightsRewriting
 
 
 class ReuseVariables(object):  # pylint: disable=useless-object-inheritance
@@ -137,49 +138,65 @@ def end(self):
 GraphRewriting.register(LayerRewriting)
 
 
-class EmbeddingLookupRewritingForVariables(EmbeddingLookupRewriting):
+class EmbeddingWeightsRewritingForVariables(EmbeddingWeightsRewriting):
   r'''Embedding lookup rewriting for variables.
   '''
   def __init__(self):
     super().__init__()
-    self._prev_lookup = None
     self._prev_get_variable = None
 
-  def build_sharded_ids(self, ids):
-    r'''Shard IDs to lookup sharded embedding weights.
-    '''
-    return ids // self.num_shards
-
-  def build_sharded_weights(self, shard, num_shards, fn, name, *args, **kwargs):
+  def build_sharded_weights(
+      self, shard, num_shards, shard_collections, fn, *args, **kwargs):
     r'''Build sharded embedding weights.
     '''
     shape = kwargs.get('shape', None)
     shape = tensor_shape.as_shape(shape)
-    if shape.dims is None:
-      return fn(name, *args, **kwargs)
-
+    shape_dims = shape.dims
     shape = shape.as_list()
-    if len(shape) != 2:
-      return fn(name, *args, **kwargs)
+    if shape_dims is None or len(shape) != 2:
+      not_embedding_weights = fn(*args, **kwargs)
+      logging.vlog(
+        1,
+        f'Embedding weights {not_embedding_weights.name.split(":")[0]} '
+        'is skipped')
+      return not_embedding_weights
 
     bucket_size, dimension_size = shape
+    batch_size = Context.get().options.batch_size
+    if bucket_size <= num_shards or bucket_size <= batch_size:
+      collections = kwargs.pop('collections', None)
+      if not collections:
+        collections = [ops.GraphKeys.GLOBAL_VARIABLES]
+      collections.append(GraphKeys.TRAINABLE_REPLICATED_SMALL)
+      kwargs['collections'] = list(set(collections))
+      embedding_weights = fn(*args, **kwargs)
+      logging.vlog(
+        1,
+        f'Embedding weights {embedding_weights.name.split(":")[0]} is small')
+      return embedding_weights
+
+    kwargs['collections'] = shard_collections
     sharded_bucket_size = bucket_size // num_shards
     if shard < bucket_size % num_shards:
       sharded_bucket_size += 1
     shape = (sharded_bucket_size, dimension_size)
     kwargs['shape'] = shape
 
-    embedding_weights = fn(f'{name}/part_{shard}', *args, **kwargs)
+    var_scope, var_store, name, *next_args = args
+    embedding_weights = fn(
+      var_scope, var_store, f'{name}/part_{shard}', *next_args, **kwargs)
     bucket_offset = (bucket_size // num_shards) * shard
     remained_buckets = bucket_size % num_shards
     if shard < remained_buckets:
       bucket_offset += shard
     else:
       bucket_offset += remained_buckets
+    full_name = embedding_weights.name.split(':')[0]
+    full_name = full_name.split('/part')[0]
     if hasattr(embedding_weights, '_set_save_slice_info'):
       embedding_weights._set_save_slice_info(  # pylint: disable=protected-access
         variables.Variable.SaveSliceInfo(
-          full_name=name,
+          full_name=full_name,
           full_shape=[bucket_size, dimension_size],
           var_offset=[bucket_offset, 0],
           var_shape=shape))
@@ -188,39 +205,27 @@ def build_sharded_weights(self, shard, num_shards, fn, name, *args, **kwargs):
         poffset = embedding_partition._get_save_slice_info().var_offset  # pylint: disable=protected-access
         embedding_partition._set_save_slice_info(  # pylint: disable=protected-access
           variables.Variable.SaveSliceInfo(
-            full_name=name,
+            full_name=full_name,
             full_shape=[bucket_size, dimension_size],
             var_offset=[bucket_offset + poffset[0], poffset[1]],
             var_shape=embedding_partition.shape))
     else:
-      logging.warning(f'Embedding weights {name} cannot be saved correctly')
+      logging.warning(
+        f'Embedding weights {full_name} cannot be saved correctly')
 
     return embedding_weights
 
   def begin(self):
     r'''Rewrites API.
     '''
-    import tensorflow as tf  # pylint: disable=import-outside-toplevel
-    self._prev_lookup = embedding_ops.embedding_lookup
-    embedding_ops.embedding_lookup = self.wraps_embedding_lookup(
-      embedding_ops.embedding_lookup)
-    tf.nn.embedding_lookup = self.wraps_embedding_lookup(
-      embedding_ops.embedding_lookup)
-
-    self._prev_get_variable = vs.get_variable
-    vs.get_variable = self.wraps_build_embedding_weights(
-      self._prev_get_variable)
-    tf.get_variable = self.wraps_build_embedding_weights(
+    self._prev_get_variable = vs.VariableScope.get_variable  # pylint: disable=protected-access
+    vs.VariableScope.get_variable = self.wraps_build_embedding_weights(  # pylint: disable=protected-access
       self._prev_get_variable)
 
   def end(self):
     r'''Revert API rewriting.
     '''
-    import tensorflow as tf  # pylint: disable=import-outside-toplevel
-    vs.get_variable = self._prev_get_variable
-    tf.get_variable = self._prev_get_variable
-    embedding_ops.embedding_lookup = self._prev_lookup
-    tf.nn.embedding_lookup = self._prev_lookup
+    vs.VariableScope.get_variable = self._prev_get_variable  # pylint: disable=protected-access
 
 
-EmbeddingLookupRewriting.register(EmbeddingLookupRewritingForVariables)
+EmbeddingWeightsRewriting.register(EmbeddingWeightsRewritingForVariables)
diff --git a/setup.py b/setup.py
index 9f54d858..09d35551 100644
--- a/setup.py
+++ b/setup.py
@@ -20,14 +20,11 @@
 from __future__ import division
 from __future__ import print_function
 
-import fnmatch
 import os
 
-from Cython.Build import cythonize
 from Cython.Distutils import build_ext
 from setuptools import find_packages
 from setuptools import setup
-from setuptools.command.build_py import build_py
 from setuptools.dist import Distribution
 from setuptools.extension import Extension
 
@@ -37,32 +34,18 @@
 NAME = f'hybridbackend{os.getenv("WHEEL_ALIAS", "")}'
 VERSION = f'{__version__}{os.getenv("WHEEL_BUILD", "")}'
 REQUIRES = os.getenv('WHEEL_REQUIRES', '').split(';')
-DEBUG = os.getenv('WHEEL_DEBUG', 'OFF')
 PACKAGES = find_packages(exclude=['tests'])
 PACKAGE_DATA = {'': ['*.so', '*.so.*']}
 
 
-class build_nonlib_py(build_py):  # pylint: disable=invalid-name
-  def find_package_modules(self, package, package_dir):
-    modules = super().find_package_modules(package, package_dir)
-    if DEBUG != 'ON':
-      modules = [
-        (pkg, mod, f) for (pkg, mod, f) in modules
-        if not fnmatch.fnmatchcase(f, pat='*/*_lib.py')]
-    return modules
-
-
-if DEBUG == 'ON':
-  EXT_SOURCES = []
-else:
-  EXT_SOURCES = ['**/*_lib.py']
-EXT_MODULES = cythonize(
-  [Extension('hybridbackend.*', sources=EXT_SOURCES)],
-  language_level=3,
-  compiler_directives=dict(always_allow_keywords=True))
-CMDCLASS = dict(build_ext=build_ext, build_py=build_nonlib_py)
+class NoExtensionBuilder(build_ext):
+  r'''Build extensions to do nothing.
+  '''
+  def build_extension(self, ext):
+    return
 
 
+CMDCLASS = dict(build_ext=NoExtensionBuilder)
 DESCRIPTION = ('A high-performance framework for training wide-and-deep '
                'recommender systems on heterogeneous cluster')
 LONG_DESCRIPTION = DESCRIPTION
@@ -89,7 +72,7 @@ def is_pure(self):
   include_package_data=True,
   package_data=PACKAGE_DATA,
   install_requires=REQUIRES,
-  ext_modules=EXT_MODULES,
+  ext_modules=[Extension('', sources=[])],
   cmdclass=CMDCLASS,
   distclass=BinaryDistribution,
   zip_safe=False,


Test patch:
diff --git a/hybridbackend/tensorflow/distribute/tests/allgather_test.py b/hybridbackend/tensorflow/distribute/tests/allgather_test.py
index 353eaf1c..dc919843 100644
--- a/hybridbackend/tensorflow/distribute/tests/allgather_test.py
+++ b/hybridbackend/tensorflow/distribute/tests/allgather_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # =============================================================================
 
-r'''Tests for Allgather.
+r'''Tests for allgather collective communication.
 '''
 
 from __future__ import absolute_import
@@ -21,17 +21,98 @@
 from __future__ import print_function
 
 import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
 import unittest
 
 import numpy as np
-import tensorflow as tf
 
 import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
+
+# pylint: disable=missing-docstring,import-outside-toplevel
+
+
+# Generate different shapes for different ranks
+def _v_shape(shape, i):
+  if not shape:
+    return shape
+  shape[0] *= (i + 1)
+  return shape
+
+
+def _test_allgather(rank, world_size, shapes):
+  r'''Test Allgather.
+  '''
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  all_inputs = []
+  allgathers = []
+  allgathervs = []
+  with tf.Graph().as_default(), hb.scope():
+    for i, shape in enumerate(shapes):
+      comm_inputs = []
+      for d in range(world_size):
+        comm_inputs.append(
+          tf.get_variable(
+            f'input{i}/replicas{d}',
+            initializer=tf.random_normal(
+              shape,
+              mean=100,
+              stddev=80,
+              seed=i * len(shapes) * 100 + d)))
+      all_inputs.append(comm_inputs)
+      allgathers.append(
+        hb.distribute.allgather(comm_inputs[rank], varying_size=False))
+      allgathervs.append(hb.distribute.allgather(comm_inputs[rank]))
+    baselines = [
+      tf.concat(all_inputs[i], 0) if shape else tf.stack(all_inputs[i])
+      for i, shape in enumerate(shapes)]
+    with tf.train.MonitoredTrainingSession('') as sess:
+      results = sess.run({
+        'baselines': baselines,
+        'allgathers': allgathers,
+        'allgathervs': allgathervs})
+      results['baselines'] = [o.tolist() for o in results['baselines']]
+      results['allgathers'] = [o.tolist() for o in results['allgathers']]
+      results['allgathervs'] = [o.tolist() for o in results['allgathervs']]
+      return results
+
+
+def _test_allgatherv(rank, world_size, shapes):
+  r'''Test Allgatherv.
+  '''
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  all_inputs = []
+  allgathervs = []
+  with tf.Graph().as_default(), hb.scope():
+    for i, shape in enumerate(shapes):
+      comm_inputs = []
+      for d in range(world_size):
+        comm_inputs.append(
+          tf.get_variable(
+            f'input{i}/replicas{d}',
+            initializer=tf.random_normal(
+              _v_shape(shape, i),
+              mean=100,
+              stddev=80,
+              seed=i * len(shapes) * 100 + d)))
+      all_inputs.append(comm_inputs)
+      allgathervs.append(hb.distribute.allgather(comm_inputs[rank]))
+    baselines = [
+      tf.concat(all_inputs[i], 0) if shape else tf.stack(all_inputs[i])
+      for i, shape in enumerate(shapes)]
+    with tf.train.MonitoredTrainingSession('') as sess:
+      results = sess.run({
+        'baselines': baselines,
+        'allgathervs': allgathervs})
+      results['baselines'] = [o.tolist() for o in results['baselines']]
+      results['allgathervs'] = [o.tolist() for o in results['allgathervs']]
+      return results
 
 
-# pylint: disable=missing-docstring
 @unittest.skipUnless(
   os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
 @unittest.skipUnless(
@@ -41,110 +122,53 @@ def setUp(self):  # pylint: disable=invalid-name
     os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
     os.environ['NCCL_DEBUG'] = 'INFO'
     os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_allgather=1'
+    os.environ['TF_CPP_VMODULE'] = (
+      'nccl_comm=1,'
+      'nccl_create=1,'
+      'nccl_allgather=1,'
+      'nccl_allgatherv=1')
 
   def tearDown(self):  # pylint: disable=invalid-name
     del os.environ['TF_CPP_VMODULE']
     del os.environ['CUDA_VISIBLE_DEVICES']
 
-  def _test_allgather(self, devices, shapes):
-    hb.context.options.update(comm_pubsub_device='')
-
-    num_comms = len(shapes)
-    num_devices = len(devices)
-    all_inputs = []
-    allgathers = []
-    allgathervs = []
-
-    with tf.Graph().as_default(), hb.scope():
-      prev_comm_gathers = None
-      prev_commv_gathers = None
-      for i in xrange(num_comms):
-        comm_inputs = []
-        comm_gathers = []
-        commv_gathers = []
-        shared_name = f'comm_{i}'
-        shared_namev = f'commv_{i}'
-        for d in xrange(num_devices):
-          with tf.device(devices[d]):
-            comm_inputs.append(
-              tf.get_variable(
-                f'input{i}/part_{d}',
-                initializer=tf.random_normal(
-                  shapes[i], mean=100, stddev=80)))
-            comm = hb.distribute.Communicator.build(shared_name, devices)
-            with tf.control_dependencies(
-                [prev_comm_gathers[d]] if prev_comm_gathers else None):
-              ga1 = comm.allgather(comm_inputs[d])
-            comm_gathers.append(ga1)
-            commv = hb.distribute.Communicator.build(shared_namev, devices)
-            with tf.control_dependencies(
-                [prev_commv_gathers[d]] if prev_commv_gathers else None):
-              ga2 = commv.allgatherv(comm_inputs[d])
-            commv_gathers.append(ga2)
-        prev_comm_gathers = comm_gathers
-        prev_commv_gathers = commv_gathers
-        all_inputs.append(comm_inputs)
-        allgathers.append(comm_gathers)
-        allgathervs.append(commv_gathers)
-      baselines = [
-        tf.concat(all_inputs[i], axis=0)
-        if shapes[i] else tf.stack(all_inputs[i])
-        for i in xrange(num_comms)]
-
-      with tf.train.MonitoredTrainingSession('') as sess:
-        expects = sess.run(baselines)
-        actuals = sess.run(allgathers)
-        actuals_by_v = sess.run(allgathervs)
-      for i in xrange(num_comms):
-        for d in xrange(num_devices):
-          np.testing.assert_allclose(actuals[i][d], expects[i], atol=1e-4)
-          np.testing.assert_allclose(actuals_by_v[i][d], expects[i], atol=1e-4)
-
-  def test_one_device_one_tensor_allgather(self):
-    self._test_allgather(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0]],
-      [[16, 4]])
-
-  def test_two_devices_scalar_tensors_allgather(self):
-    self._test_allgather(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0, 1]],
-      [[]])
-
-  def test_two_devices_four_tensors_allgather(self):
-    self._test_allgather(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0, 1]],
-      [[17, 41, 37], [19, 121], [1], [16, 6]])
-
-  def test_allgather_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[2, 10])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.allgather(a)
-        loss0 = tf.reduce_mean(recv0) * 20.0
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[2, 10])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        recv1 = comm1.allgather(b * 0.75)
-        loss1 = tf.reduce_mean(recv1) * 10.0
-      loss = loss0 * loss1
-      grad0, grad1 = tf.gradients([loss], [a, b], [2.0])
-      with tf.train.MonitoredTrainingSession('') as sess:
-        g0, g1 = sess.run([grad0, grad1])
+  def _test_allgather(self, world_size, shapes):
+    results = hbtest.Spawn(world_size)(
+      lambda rank: _test_allgather(rank, world_size, shapes))
+    for d in range(world_size):
+      r = results[d]
+      for i, _ in enumerate(shapes):
         np.testing.assert_allclose(
-          g0,
-          [[25.0] * 10] * 2,
-          rtol=1e-6)
+          r['baselines'][i], r['allgathers'][i], atol=1e-4)
         np.testing.assert_allclose(
-          g1,
-          [[18.75] * 10] * 2,
-          rtol=1e-6)
+          r['baselines'][i], r['allgathervs'][i], atol=1e-4)
+
+  def test_allgather_scalar(self):
+    self._test_allgather(2, [[]])
+
+  def test_allgather_fallback(self):
+    self._test_allgather(1, [[6, 14]])
+
+  def test_allgather_n(self):
+    self._test_allgather(2, [[17, 3, 4], [2, 3], [1], [16, 6]])
+
+  def _test_allgatherv(self, world_size, shapes):
+    results = hbtest.Spawn(world_size)(
+      lambda rank: _test_allgatherv(rank, world_size, shapes))
+    for d in range(world_size):
+      r = results[d]
+      for i, _ in enumerate(shapes):
+        np.testing.assert_allclose(
+          r['baselines'][i], r['allgathervs'][i], atol=1e-4)
+
+  def test_allgatherv_scalar(self):
+    self._test_allgatherv(2, [[]])
+
+  def test_allgatherv_fallback(self):
+    self._test_allgatherv(1, [[6, 14]])
+
+  def test_allgatherv_n(self):
+    self._test_allgatherv(2, [[2, 3, 4], [2, 3], [1], [2, 6]])
 
 
 if __name__ == '__main__':
diff --git a/hybridbackend/tensorflow/distribute/tests/allgatherv_test.py b/hybridbackend/tensorflow/distribute/tests/allgatherv_test.py
deleted file mode 100644
index ca8e2203..00000000
--- a/hybridbackend/tensorflow/distribute/tests/allgatherv_test.py
+++ /dev/null
@@ -1,130 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for Allgatherv.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class AllgathervTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_allgatherv=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def _test_allgatherv(self, devices, shapes):
-    hb.context.options.update(comm_pubsub_device='')
-
-    num_comms = len(shapes)
-    num_devices = len(devices)
-    all_inputs = []
-    allgathervs = []
-
-    # Generate different shapes for different ranks
-    def _v_shape(shape, i):
-      if not shape:
-        return shape
-      shape[0] *= (i + 1)
-      return shape
-
-    with tf.Graph().as_default(), hb.scope():
-      prev_comm_gathervs = None
-      for i in xrange(num_comms):
-        comm_inputs = []
-        comm_gathervs = []
-        shared_name = f'comm_{i}'
-        for d in xrange(num_devices):
-          with tf.device(devices[d]):
-            comm_inputs.append(
-              tf.get_variable(
-                f'input{i}/part_{d}',
-                initializer=tf.random_normal(
-                  _v_shape(shapes[i], i), mean=100, stddev=80)))
-            comm = hb.distribute.Communicator.build(shared_name, devices)
-            with tf.control_dependencies(
-                [prev_comm_gathervs[d]] if prev_comm_gathervs else None):
-              comm_gathervs.append(comm.allgatherv(comm_inputs[d]))
-        prev_comm_gathervs = comm_gathervs
-        all_inputs.append(comm_inputs)
-        allgathervs.append(comm_gathervs)
-      baselines = [tf.concat(all_inputs[i], axis=0)
-                   for i in xrange(num_comms)]
-
-      with tf.train.MonitoredTrainingSession('') as sess:
-        expects = sess.run(baselines)
-        actuals = sess.run(allgathervs)
-      for i in xrange(num_comms):
-        for d in xrange(num_devices):
-          np.testing.assert_allclose(actuals[i][d], expects[i], atol=1e-4)
-
-  def test_one_device_one_tensor_allgatherv(self):
-    self._test_allgatherv(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0]],
-      [[16, 4]])
-
-  def test_two_devices_four_tensors_allgatherv(self):
-    self._test_allgatherv(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0, 1]],
-      [[17, 41, 37], [19, 121], [1], [16, 6]])
-
-  def test_allgatherv_grad_none(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default() as graph:
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[2, 10])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.allgatherv(a)
-        loss0 = tf.reduce_mean(recv0) * 20.0
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[2, 10])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        recv1 = comm1.allgatherv(b * 0.75)
-        loss1 = tf.reduce_mean(recv1) * 10.0
-      loss = loss0 * loss1
-      grad0, grad1 = tf.gradients([loss], [a, b], [2.0])
-      np.testing.assert_equal(grad0, None)
-      np.testing.assert_equal(grad1, None)
-    graph.finalize()
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/allreduce_grad_test.py b/hybridbackend/tensorflow/distribute/tests/allreduce_grad_test.py
deleted file mode 100644
index 07b3f945..00000000
--- a/hybridbackend/tensorflow/distribute/tests/allreduce_grad_test.py
+++ /dev/null
@@ -1,102 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for Allreduce gradients.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class AllreduceGradTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_allreduce=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def test_allreduce_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[2, 10])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.allreduce(a * 0.75)
-        loss0 = tf.reduce_mean(recv0) * 20.0
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[2, 10])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        recv1 = comm1.allreduce(b)
-        loss1 = tf.reduce_mean(recv1) * 10.0
-      loss = loss0 * loss1
-      grad0, grad1 = tf.gradients(
-        [loss], [a, b], [2.0], colocate_gradients_with_ops=True)
-      with tf.train.MonitoredTrainingSession() as sess:
-        g0, g1 = sess.run([grad0, grad1])
-        np.testing.assert_allclose(
-          g0,
-          [[82.5] * 10] * 2,
-          rtol=1e-6)
-        np.testing.assert_allclose(
-          g1,
-          [[110.0] * 10] * 2,
-          rtol=1e-6)
-
-  def test_allreduce_grad_error(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default():
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[2, 10])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.allreduce(a * 0.75, reduce_op=hb.distribute.ops.MAX)
-        loss0 = tf.reduce_mean(recv0) * 20.0
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[2, 10])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        recv1 = comm1.allreduce(b, reduce_op=hb.distribute.ops.MAX)
-        loss1 = tf.reduce_mean(recv1) * 10.0
-      loss = loss0 * loss1
-      with self.assertRaises(RuntimeError):
-        tf.gradients([loss], [a, b], [2.0])
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/allreduce_test.py b/hybridbackend/tensorflow/distribute/tests/allreduce_test.py
index 5699b122..b76cd0a0 100644
--- a/hybridbackend/tensorflow/distribute/tests/allreduce_test.py
+++ b/hybridbackend/tensorflow/distribute/tests/allreduce_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # =============================================================================
 
-r'''Tests for Allreduce.
+r'''Tests for allreduce collective communication.
 '''
 
 from __future__ import absolute_import
@@ -21,61 +21,72 @@
 from __future__ import print_function
 
 import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
 import unittest
 
 import numpy as np
 
 import hybridbackend.common.test as hbtest
 
+# pylint: disable=missing-docstring,import-outside-toplevel
 
-# pylint: disable=missing-docstring
-def _test_simple_allreduce(rank, a, b):
-  # pylint: disable=import-outside-toplevel
+
+def _test_allreduce(rank, a, b):
   import tensorflow as tf
 
   import hybridbackend.tensorflow as hb
 
   with tf.Graph().as_default():
     with hb.scope():
-      comm0 = hb.distribute.Communicator.build('comm0', hb.context.devices)
       input0 = tf.constant(a) if rank == 0 else tf.constant(b)
-      sum0 = comm0.allreduce(input0)
+      sum0 = hb.distribute.allreduce(input0)
       with tf.train.MonitoredTrainingSession('') as sess:
         return sess.run(sum0)
 
 
-def _test_simple_allreduce_max(rank, a, b):
-  # pylint: disable=import-outside-toplevel
+def _test_allreduce_max(rank, a, b):
   import tensorflow as tf
 
   import hybridbackend.tensorflow as hb
 
   with tf.Graph().as_default():
     with hb.scope():
-      comm0 = hb.distribute.Communicator.build('comm0', hb.context.devices)
       input0 = tf.constant(a) if rank == 0 else tf.constant(b)
-      sum0 = comm0.allreduce(input0, reduce_op=hb.distribute.ops.MAX)
+      sum0 = hb.distribute.allreduce(
+        input0, reduce_op=hb.distribute.ops.MAX)
       with tf.train.MonitoredTrainingSession('') as sess:
         return sess.run(sum0)
 
 
-def _test_simple_allreduce_multicomm(rank, a, b, ncomms):
-  # pylint: disable=import-outside-toplevel
+def _test_allreduce_n(rank, a, b):
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  with tf.Graph().as_default():
+    with hb.scope():
+      input0 = tf.constant(a) if rank == 0 else tf.constant(b)
+      sum0 = hb.distribute.allreduce(input0)
+      sum1 = hb.distribute.allreduce(input0)
+      sum2 = hb.distribute.allreduce(sum0)
+      result = sum1 + sum2
+      with tf.train.MonitoredTrainingSession('') as sess:
+        return sess.run(result)
+
+
+def _test_allreduce_grad(rank, g, shp, a, b):
   import tensorflow as tf
 
   import hybridbackend.tensorflow as hb
 
   with tf.Graph().as_default():
     with hb.scope():
-      sums = []
-      for icomm in xrange(ncomms):
-        commi = hb.distribute.Communicator.build(
-          f'comm{icomm}', hb.context.devices)
-        inputi = tf.constant(a) if rank == 0 else tf.constant(b + icomm)
-        sums.append(commi.allreduce(inputi))
+      data = tf.constant(a if rank == 0 else b, shape=shp)
+      reduced = hb.distribute.allreduce(data)
+      loss = tf.reduce_mean(reduced)
+      grad = tf.gradients(
+        [loss], [data], [g], colocate_gradients_with_ops=True)
       with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(sums)
+        return sess.run(grad)
 
 
 @unittest.skipUnless(
@@ -87,43 +98,55 @@ def setUp(self):  # pylint: disable=invalid-name
     os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
     os.environ['NCCL_DEBUG'] = 'INFO'
     os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_allreduce=1'
+    os.environ['TF_CPP_VMODULE'] = (
+      'nccl_comm=1,'
+      'nccl_create=1,'
+      'nccl_allreduce=1')
 
   def tearDown(self):  # pylint: disable=invalid-name
     del os.environ['TF_CPP_VMODULE']
     del os.environ['CUDA_VISIBLE_DEVICES']
 
-  def test_simple_allreduce(self):
+  def test_allreduce(self):
     a = 13
     b = 22
     results = hbtest.Spawn(2)(
-      lambda rank: _test_simple_allreduce(rank, a, b))
+      lambda rank: _test_allreduce(rank, a, b))
     np.testing.assert_allclose(results[0], a + b, rtol=1e-6)
     np.testing.assert_allclose(results[1], a + b, rtol=1e-6)
 
-  def test_onedevice_allreduce(self):
+  def test_allreduce_fallback(self):
     a = 13
     results = hbtest.Spawn(1)(
-      lambda rank: _test_simple_allreduce(rank, a, 22))
+      lambda rank: _test_allreduce(rank, a, 22))
     np.testing.assert_allclose(results[0], a, rtol=1e-6)
 
-  def test_allreduce_max(self):
+  def xtest_allreduce_max(self):
     a = 13
     b = 22
     results = hbtest.Spawn(2)(
-      lambda rank: _test_simple_allreduce_max(rank, a, b))
+      lambda rank: _test_allreduce_max(rank, a, b))
     np.testing.assert_allclose(results[0], b, rtol=1e-6)
     np.testing.assert_allclose(results[1], b, rtol=1e-6)
 
-  def test_multicomm(self):
+  def test_allreduce_n(self):
     a = 13
     b = 22
-    ncomms = 10
     results = hbtest.Spawn(2)(
-      lambda rank: _test_simple_allreduce_multicomm(rank, a, b, ncomms))
-    for icomm in xrange(ncomms):
-      np.testing.assert_allclose(results[0][icomm], a + b + icomm, rtol=1e-6)
-      np.testing.assert_allclose(results[1][icomm], a + b + icomm, rtol=1e-6)
+      lambda rank: _test_allreduce_n(rank, a, b))
+    np.testing.assert_allclose(results[0], 3 * (a + b), rtol=1e-6)
+    np.testing.assert_allclose(results[1], 3 * (a + b), rtol=1e-6)
+
+  def test_allreduce_grad(self):
+    g = 2.0
+    w = 4
+    h = 8
+    results = hbtest.Spawn(2)(
+      lambda rank: _test_allreduce_grad(rank, g, [w, h], 1.0, 2.0))
+    np.testing.assert_allclose(results[0], results[1], rtol=1e-6)
+    np.testing.assert_allclose(
+      results[0],
+      [[[g / (w * h) * 2] * h] * w], rtol=1e-6)
 
 
 if __name__ == '__main__':
diff --git a/hybridbackend/tensorflow/distribute/tests/alltoall_test.py b/hybridbackend/tensorflow/distribute/tests/alltoall_test.py
index c4e91e06..d6ea18b7 100644
--- a/hybridbackend/tensorflow/distribute/tests/alltoall_test.py
+++ b/hybridbackend/tensorflow/distribute/tests/alltoall_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # =============================================================================
 
-r'''Tests for Alltoall.
+r'''Tests for alltoall collective communication.
 '''
 
 from __future__ import absolute_import
@@ -24,13 +24,158 @@
 import unittest
 
 import numpy as np
-import tensorflow as tf
 
 import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
+
+# pylint: disable=missing-docstring,import-outside-toplevel
+
+
+def _test_alltoall(rank, world_size, shape):
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  with tf.Graph().as_default(), hb.scope():
+    inputs = []
+    for d in range(world_size):
+      inputs.append(
+        tf.get_variable(
+          f'alltoall_input/replica{d}',
+          dtype=tf.float32,
+          initializer=tf.random_normal(shape, mean=100, stddev=80, seed=d)))
+    actual = hb.distribute.alltoall(inputs[rank])
+
+    with tf.train.MonitoredTrainingSession('') as sess:
+      results = sess.run({'actual': actual, 'inputs': inputs})
+      results['actual'] = [o.tolist() for o in results['actual']]
+      results['inputs'] = [o.tolist() for o in results['inputs']]
+      results['expected'] = list(map(list, zip(*results['inputs'])))
+      return results
+
+
+def _test_alltoall_grad(rank, world_size, shape, g):
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  with tf.Graph().as_default(), hb.scope():
+    inputs = []
+    for d in range(world_size):
+      inputs.append(
+        tf.get_variable(
+          f'alltoall_input/replica{d}',
+          dtype=tf.float32,
+          initializer=tf.random_normal(shape, mean=100, stddev=80, seed=d)))
+    shuffled = hb.distribute.alltoall(inputs[rank])
+    loss = tf.reduce_mean(shuffled)
+    grad = tf.gradients(
+      [loss], [inputs[rank]], [g],
+      colocate_gradients_with_ops=True)
+
+    with tf.train.MonitoredTrainingSession('') as sess:
+      return sess.run(grad)
+
+
+def _test_alltoallv(rank, ids_list, sizes_list, dtype=None, wire_dtype=None):
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  with tf.Graph().as_default(), hb.scope(comm_wire_dtype=wire_dtype):
+    ids = tf.constant(
+      ids_list[rank], dtype=tf.int64 if dtype is None else dtype)
+    sizes = tf.constant(sizes_list[rank], dtype=tf.int32)
+    out_ids, out_sizes = hb.distribute.alltoall(ids, sizes=sizes)
+
+    with tf.train.MonitoredTrainingSession('') as sess:
+      return sess.run({'ids': out_ids, 'sizes': out_sizes})
+
+
+def _test_alltoallv_fp16(rank, ids_list, sizes_list):
+  import tensorflow as tf
+
+  return _test_alltoallv(
+    rank, ids_list, sizes_list, dtype=tf.float32, wire_dtype=tf.float16)
+
+
+def _test_alltoallv_grad(rank, values, sizes, g):
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  with tf.Graph().as_default(), hb.scope():
+    inputs = tf.constant(
+      values[rank],
+      dtype=tf.float32,
+      shape=(sum(sizes[rank]),))
+    input_sizes = tf.constant(
+      sizes[rank],
+      dtype=tf.int32)
+    outputs, _ = hb.distribute.alltoall(inputs, sizes=input_sizes)
+    loss = tf.reduce_mean(outputs)
+    grad = tf.gradients(
+      [loss], [inputs], [g],
+      colocate_gradients_with_ops=True)
+
+    with tf.train.MonitoredTrainingSession('') as sess:
+      return sess.run(grad)
+
+
+def _test_alltoallv_n(rank, ids_and_sizes_list, wire_dtype=None):
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  with tf.Graph().as_default(), hb.scope(comm_wire_dtype=wire_dtype):
+    out_ids_list = []
+    out_sizes_list = []
+    for ids_and_sizes in ids_and_sizes_list[rank]:
+      ids = tf.constant(ids_and_sizes['ids'], dtype=tf.float32)
+      sizes = tf.constant(ids_and_sizes['sizes'], dtype=tf.int32)
+      out_ids, out_sizes = hb.distribute.alltoall(ids, sizes=sizes)
+      out_ids_list.append(out_ids)
+      out_sizes_list.append(out_sizes)
+
+    with tf.train.MonitoredTrainingSession('') as sess:
+      results = sess.run({'ids': out_ids_list, 'sizes': out_sizes_list})
+      return [
+        {'ids': ids.tolist(), 'sizes': results['sizes'][i].tolist()}
+        for i, ids in enumerate(results['ids'])]
+
+
+def _test_alltoallv_n_fp16(rank, ids_and_sizes_list):
+  import tensorflow as tf
+
+  return _test_alltoallv_n(rank, ids_and_sizes_list, wire_dtype=tf.float16)
+
+
+def _test_alltoallv_n_grad(rank, ids_and_sizes_list, g):
+  import tensorflow as tf
+
+  import hybridbackend.tensorflow as hb
+
+  with tf.Graph().as_default(), hb.scope():
+    grads = []
+    for ids_and_sizes in ids_and_sizes_list[rank]:
+      inputs = tf.constant(
+        ids_and_sizes['ids'],
+        dtype=tf.float32,
+        shape=(sum(ids_and_sizes['sizes']),))
+      input_sizes = tf.constant(
+        ids_and_sizes['sizes'],
+        dtype=tf.int32)
+      outputs, _ = hb.distribute.alltoall(inputs, sizes=input_sizes)
+      loss = tf.reduce_mean(outputs)
+      grads.append(
+        tf.gradients(
+          [loss], [inputs], [g],
+          colocate_gradients_with_ops=True))
+
+    with tf.train.MonitoredTrainingSession('') as sess:
+      results = sess.run(grads)
+      return [r[0].tolist() for r in results]
 
 
-# pylint: disable=missing-docstring
 @unittest.skipUnless(
   os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
 @unittest.skipUnless(
@@ -40,83 +185,123 @@ def setUp(self):  # pylint: disable=invalid-name
     os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
     os.environ['NCCL_DEBUG'] = 'INFO'
     os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_alltoall=1'
+    os.environ['TF_CPP_VMODULE'] = (
+      'nccl_comm=1,'
+      'nccl_create=1,'
+      'nccl_alltoall=1,'
+      'nccl_alltoallv=1,'
+      'packing=2,'
+      'optimize_collective=2')
 
   def tearDown(self):  # pylint: disable=invalid-name
     del os.environ['TF_CPP_VMODULE']
     del os.environ['CUDA_VISIBLE_DEVICES']
 
-  def test_alltoall(self):
-    hb.context.options.update(comm_pubsub_device='')
+  def xtest_alltoall(self):
+    world_size = 2
+    results = hbtest.Spawn(world_size)(
+      lambda rank: _test_alltoall(rank, world_size, [2, 3]))
+    for d in range(world_size):
+      np.testing.assert_allclose(
+        results[d]['actual'], results[d]['expected'][d], rtol=1e-6)
 
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoall_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    shape = [2, 4]
-    inputs = []
-    expected = []
-    actual = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            initializer = tf.random_normal(shape, mean=100, stddev=80)
-            inputs.append(tf.get_variable(
-              f'alltoall_input_{i}',
-              initializer=initializer))
-            transposed = comm.alltoall(inputs[i])
-            actual.append(transposed)
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'inputs': inputs})
-          expected = list(map(list, zip(*results['inputs'])))
-          np.testing.assert_allclose(expected, results['actual'], rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoall_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoall_grad_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [9.6, 8.8]
-    shape = [2, 10]
-    grad_ys = [2.0]
-    all_ys = []
-    all_outputs = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            inputs = tf.constant(values[i], shape=shape)
-            outputs = comm.alltoall(inputs)
-            all_inputs.append(inputs)
-            all_outputs.append(outputs)
-            all_ys.append(tf.reduce_sum(outputs, axis=0))
-
-        with tf.device(devices[0]):
-          xs = all_inputs
-          ys = [tf.reduce_sum(tf.add_n(all_ys))]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-
-          all_inputs_v2 = [tf.split(x, 2) for x in all_inputs]
-          baseline_all_outputs = list(map(list, zip(*all_inputs_v2)))
-          baseline_all_ys = [
-            tf.reduce_sum(baseline_all_outputs, axis=0)]
-          baseline_ys = [tf.reduce_sum(tf.add_n(baseline_all_ys))]
-          expected = tf.gradients(
-            baseline_ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          np.testing.assert_allclose(
-            results['expected'], results['actual'], rtol=1e-6)
-    finally:
-      del server
+  def xtest_alltoall_grad(self):
+    world_size = 2
+    w = 2
+    h = 10
+    g = 2.0
+    results = hbtest.Spawn(world_size)(
+      lambda rank: _test_alltoall_grad(rank, world_size, [w, h], g))
+    np.testing.assert_allclose(results[0], results[1], rtol=1e-6)
+    np.testing.assert_allclose(
+      results[0], [[[g / (w * h)] * h] * w], rtol=1e-6)
+
+  def xtest_alltoallv(self):
+    results = hbtest.Spawn(2)(
+      lambda rank: _test_alltoallv(
+        rank, [[1, 2, 3], [4, 5, 6]], [[1, 2], [1, 2]]))
+    np.testing.assert_allclose(results[0]['ids'], [1, 4], rtol=1e-6)
+    np.testing.assert_allclose(results[0]['sizes'], [1, 1], rtol=1e-6)
+    np.testing.assert_allclose(results[1]['ids'], [2, 3, 5, 6], rtol=1e-6)
+    np.testing.assert_allclose(results[1]['sizes'], [2, 2], rtol=1e-6)
+
+  def xtest_alltoallv_grad(self):
+    g = 2.0
+    sizes = [[5, 1], [3, 4]]
+    results = hbtest.Spawn(2)(
+      lambda rank: _test_alltoallv_grad(
+        rank, [3.6, 4.2], sizes, g))
+    g0 = g / (sizes[0][0] + sizes[1][0])
+    g1 = g / (sizes[0][1] + sizes[1][1])
+    np.testing.assert_allclose(
+      results[0],
+      [sizes[0][0] * [g0] + sizes[0][1] * [g1]],
+      rtol=1e-6)
+    np.testing.assert_allclose(
+      results[1],
+      [sizes[1][0] * [g0] + sizes[1][1] * [g1]],
+      rtol=1e-6)
+
+  def xtest_alltoallv_fp16(self):
+    results = hbtest.Spawn(2)(
+      lambda rank: _test_alltoallv_fp16(
+        rank, [[1., 2., 3.], [4., 5., 6.]], [[1, 2], [1, 2]]))
+    np.testing.assert_allclose(results[0]['ids'], [1., 4.], rtol=1e-6)
+    np.testing.assert_allclose(results[0]['sizes'], [1, 1], rtol=1e-6)
+    np.testing.assert_allclose(results[1]['ids'], [2., 3., 5., 6.], rtol=1e-6)
+    np.testing.assert_allclose(results[1]['sizes'], [2, 2], rtol=1e-6)
+
+  def test_alltoallv_n(self):
+    results = hbtest.Spawn(2)(
+      lambda rank: _test_alltoallv_n(
+        rank,
+        {0: [{'ids': [1., 2., 3.], 'sizes': [1, 2]},
+             {'ids': [4., 5., 6.], 'sizes': [2, 1]}],
+         1: [{'ids': [7., 8., 9.], 'sizes': [2, 1]},
+             {'ids': [10., 11., 12.], 'sizes': [1, 2]}]}))
+    np.testing.assert_allclose(results[0][0]['ids'], [1., 7., 8.], rtol=1e-6)
+    np.testing.assert_equal(results[0][0]['sizes'], [1, 2])
+    np.testing.assert_allclose(results[1][0]['ids'], [2., 3., 9.], rtol=1e-6)
+    np.testing.assert_equal(results[1][0]['sizes'], [2, 1])
+    np.testing.assert_allclose(results[0][1]['ids'], [4., 5., 10.], rtol=1e-6)
+    np.testing.assert_equal(results[0][1]['sizes'], [2, 1])
+    np.testing.assert_allclose(results[1][1]['ids'], [6., 11., 12.], rtol=1e-6)
+    np.testing.assert_equal(results[1][1]['sizes'], [1, 2])
+
+  def xtest_alltoallv_n_fp16(self):
+    results = hbtest.Spawn(2)(
+      lambda rank: _test_alltoallv_n_fp16(
+        rank,
+        {0: [{'ids': [1., 2., 3.], 'sizes': [1, 2]},
+             {'ids': [4., 5., 6.], 'sizes': [2, 1]}],
+         1: [{'ids': [7., 8., 9.], 'sizes': [2, 1]},
+             {'ids': [10., 11., 12.], 'sizes': [1, 2]}]}))
+    np.testing.assert_allclose(results[0][0]['ids'], [1., 7., 8.], rtol=1e-6)
+    np.testing.assert_equal(results[0][0]['sizes'], [1, 2])
+    np.testing.assert_allclose(results[1][0]['ids'], [2., 3., 9.], rtol=1e-6)
+    np.testing.assert_equal(results[1][0]['sizes'], [2, 1])
+    np.testing.assert_allclose(results[0][1]['ids'], [4., 5., 10.], rtol=1e-6)
+    np.testing.assert_equal(results[0][1]['sizes'], [2, 1])
+    np.testing.assert_allclose(results[1][1]['ids'], [6., 11., 12.], rtol=1e-6)
+    np.testing.assert_equal(results[1][1]['sizes'], [1, 2])
+
+  def xtest_alltoallv_n_grad(self):
+    results = hbtest.Spawn(2)(
+      lambda rank: _test_alltoallv_n_grad(
+        rank,
+        {0: [{'ids': [1., 2., 3.], 'sizes': [1, 2]},
+             {'ids': [4., 5., 6.], 'sizes': [2, 1]}],
+         1: [{'ids': [7., 8., 9.], 'sizes': [2, 1]},
+             {'ids': [10., 11., 12.], 'sizes': [1, 2]}]},
+        2.0))
+    np.testing.assert_allclose(
+      results[0][0], [0.666667, 0.666667, 0.666667], rtol=1e-6)
+    np.testing.assert_allclose(
+      results[0][1], [0.666667, 0.666667, 0.666667], rtol=1e-6)
+    np.testing.assert_allclose(
+      results[1][0], [0.666667, 0.666667, 0.666667], rtol=1e-6)
+    np.testing.assert_allclose(
+      results[1][1], [0.666667, 0.666667, 0.666667], rtol=1e-6)
 
 
 if __name__ == '__main__':
diff --git a/hybridbackend/tensorflow/distribute/tests/alltoallv_n_test.py b/hybridbackend/tensorflow/distribute/tests/alltoallv_n_test.py
deleted file mode 100644
index f01a219d..00000000
--- a/hybridbackend/tensorflow/distribute/tests/alltoallv_n_test.py
+++ /dev/null
@@ -1,177 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for AlltoallvN.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class AlltoallvNTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_alltoallv=1,nccl_alltoallv_n=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def test_alltoallv_n(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallv_n_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    try:
-      with tf.Graph().as_default():
-        train_ops = []
-        with tf.device(devices[0]):
-          comm0 = hb.distribute.Communicator.build(comm_id, devices)
-          ids0 = tf.constant([1, 2, 3], dtype=tf.int64)
-          sizes0 = tf.constant([1, 2], dtype=tf.int32)
-          ids2 = tf.constant([7, 8, 9], dtype=tf.int64)
-          sizes2 = tf.constant([2, 1], dtype=tf.int32)
-          out02, out_sizes02 = comm0.alltoallv_n(
-            [ids0, ids2], [sizes0, sizes2])
-          train_ops.append([out02[0], out_sizes02[0], out02[1], out_sizes02[1]])
-        with tf.device(devices[1]):
-          comm1 = hb.distribute.Communicator.build(comm_id, devices)
-          ids1 = tf.constant([4, 5, 6], dtype=tf.int64)
-          sizes1 = tf.constant([2, 1], dtype=tf.int32)
-          ids3 = tf.constant([3, 9, 6], dtype=tf.int64)
-          sizes3 = tf.constant([1, 2], dtype=tf.int32)
-          out13, out_sizes13 = comm1.alltoallv_n(
-            [ids1, ids3], [sizes1, sizes3])
-          train_ops.append([out13[0], out_sizes13[0], out13[1], out_sizes13[1]])
-        with server.monitored_session() as sess:
-          d0result, d1result = sess.run(train_ops)
-          np.testing.assert_equal(d0result[0], [1, 4, 5])
-          np.testing.assert_equal(d0result[1], [1, 2])
-          np.testing.assert_equal(d0result[2], [7, 8, 3])
-          np.testing.assert_equal(d0result[3], [2, 1])
-          np.testing.assert_equal(d1result[0], [2, 3, 6])
-          np.testing.assert_equal(d1result[1], [2, 1])
-          np.testing.assert_equal(d1result[2], [9, 9, 6])
-          np.testing.assert_equal(d1result[3], [1, 2])
-    finally:
-      del server
-
-  def test_alltoallv_n_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallv_n_grad_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [[1.0, 2.4], [9.6, 8.8]]
-    shapes = [[[6], [8]], [[10], [12]]]
-    lrs = [1.0, 3.0]
-    grad_ys = [2.0]
-    all_ys = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            inputs = [
-              tf.constant(values[i][j], shape=shapes[i][j])
-              for j, _ in enumerate(devices)]
-            inputs = tf.concat(inputs, axis=0)
-            all_inputs.append(inputs)
-            sizes = tf.constant(
-              [shapes[i][j][0] for j, _ in enumerate(devices)])
-            outputs, _ = comm.alltoallv_n(
-              [inputs], [sizes])
-            all_ys.append(lrs[i] * tf.reduce_sum(outputs[0]))
-        with tf.device(devices[0]):
-          xs = all_inputs
-          ys = [tf.add_n(all_ys)]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          actual = sess.run(actual)
-          np.testing.assert_allclose(
-            actual[0],
-            [2., 2., 2., 2., 2., 2.,
-             6., 6., 6., 6., 6., 6., 6., 6.],
-            rtol=1e-6)
-          np.testing.assert_allclose(
-            actual[1],
-            [2., 2., 2., 2., 2., 2., 2., 2., 2., 2.,
-             6., 6., 6., 6., 6., 6., 6., 6., 6., 6., 6., 6.],
-            rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoallv_n_grad_half_prec(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallv_n_grad_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [[1.0, 2.4], [9.6, 8.8]]
-    shapes = [[[6], [8]], [[10], [12]]]
-    lrs = [1.0, 3.0]
-    grad_ys = [2.0]
-    all_ys = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            inputs = [
-              tf.constant(values[i][j], shape=shapes[i][j])
-              for j, _ in enumerate(devices)]
-            inputs = tf.concat(inputs, axis=0)
-            all_inputs.append(inputs)
-            sizes = tf.constant(
-              [shapes[i][j][0] for j, _ in enumerate(devices)])
-            outputs, _ = comm.alltoallv_n(
-              [inputs], [sizes],
-              wire_dtype=tf.float16)
-            all_ys.append(lrs[i] * tf.reduce_sum(outputs[0]))
-        with tf.device(devices[0]):
-          xs = all_inputs
-          ys = [tf.add_n(all_ys)]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          sess.run(actual)
-    finally:
-      del server
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/alltoallv_test.py b/hybridbackend/tensorflow/distribute/tests/alltoallv_test.py
deleted file mode 100644
index 4cb1ba26..00000000
--- a/hybridbackend/tensorflow/distribute/tests/alltoallv_test.py
+++ /dev/null
@@ -1,303 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for Alltoallv.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class AlltoallvTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_alltoallv=1,nccl_alltoallv_n=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def test_alltoallv(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallv_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    try:
-      with tf.Graph().as_default():
-        train_ops = []
-        with tf.device(devices[0]):
-          comm0 = hb.distribute.Communicator.build(comm_id, devices)
-          ids0 = tf.constant([1, 2, 3], dtype=tf.int64)
-          sizes0 = tf.constant([1, 2], dtype=tf.int32)
-          out0, out_sizes0 = comm0.alltoallv(ids0, sizes0)
-          train_ops.append([out0, out_sizes0])
-        with tf.device(devices[1]):
-          comm1 = hb.distribute.Communicator.build(comm_id, devices)
-          ids1 = tf.constant([4, 5, 6], dtype=tf.int64)
-          sizes1 = tf.constant([1, 2], dtype=tf.int32)
-          out1, out_sizes1 = comm1.alltoallv(ids1, sizes1)
-          train_ops.append([out1, out_sizes1])
-        with server.monitored_session() as sess:
-          d0result, d1result = sess.run(train_ops)
-          np.testing.assert_allclose(d0result[0], [1, 4], rtol=1e-6)
-          np.testing.assert_allclose(d0result[1], [1, 1], rtol=1e-6)
-          np.testing.assert_allclose(d1result[0], [2, 3, 5, 6], rtol=1e-6)
-          np.testing.assert_allclose(d1result[1], [2, 2], rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoallv_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallv_grad_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [[1.0, 2.4], [9.6, 8.8]]
-    shapes = [[[6], [8]], [[10], [12]]]
-    lrs = [1.0, 3.0]
-    grad_ys = [2.0]
-    all_ys = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            inputs = [
-              tf.constant(values[i][j], shape=shapes[i][j])
-              for j, _ in enumerate(devices)]
-            inputs = tf.concat(inputs, axis=0)
-            all_inputs.append(inputs)
-            sizes = tf.constant(
-              [shapes[i][j][0] for j, _ in enumerate(devices)])
-            outputs, _ = comm.alltoallv(inputs, sizes)
-            all_ys.append(lrs[i] * tf.reduce_sum(outputs))
-        with tf.device(devices[0]):
-          xs = all_inputs
-          ys = [tf.add_n(all_ys)]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          actual = sess.run(actual)
-          np.testing.assert_allclose(
-            actual[0],
-            [2., 2., 2., 2., 2., 2.,
-             6., 6., 6., 6., 6., 6., 6., 6.],
-            rtol=1e-6)
-          np.testing.assert_allclose(
-            actual[1],
-            [2., 2., 2., 2., 2., 2., 2., 2., 2., 2.,
-             6., 6., 6., 6., 6., 6., 6., 6., 6., 6., 6., 6.],
-            rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoallv_half_prec(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    os.environ['HB_COMM_WIRE_DTYPE_FLOAT'] = 'float16'
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallv_half_prec_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    try:
-      with tf.Graph().as_default():
-        train_ops = []
-        with tf.device(devices[0]):
-          comm0 = hb.distribute.Communicator.build(comm_id, devices)
-          ids0 = tf.constant(
-            [1.0123, 2.3333, 3.4444], dtype=tf.float32)
-          sizes0 = tf.constant([1, 2], dtype=tf.int32)
-          out0, out_sizes0 = comm0.alltoallv(ids0, sizes0)
-          train_ops.append([out0, out_sizes0])
-        with tf.device(devices[1]):
-          comm1 = hb.distribute.Communicator.build(comm_id, devices)
-          ids1 = tf.constant(
-            [4.3333, 5.4444, 6.5555], dtype=tf.float32)
-          sizes1 = tf.constant([1, 2], dtype=tf.int32)
-          out1, out_sizes1 = comm1.alltoallv(ids1, sizes1)
-          train_ops.append([out1, out_sizes1])
-        with server.monitored_session() as sess:
-          d0result, d1result = sess.run(train_ops)
-          np.testing.assert_allclose(d0result[0], [1.0123, 4.3333], rtol=1e-2)
-          np.testing.assert_allclose(d0result[1], [1, 1], rtol=1e-6)
-          np.testing.assert_allclose(
-            d1result[0], [2.3333, 3.4444, 5.4444, 6.5555], rtol=1e-2)
-          np.testing.assert_allclose(d1result[1], [2, 2], rtol=1e-6)
-    finally:
-      del server
-    del os.environ['HB_COMM_WIRE_DTYPE_FLOAT']
-
-  def test_alltoallv_multi_steps(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallv_multistep_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    shapes = [[10230 + 5260, 64], [4000 + 990, 64]]
-    sizes = [[10230, 5260], [4000, 990]]
-    train_ops = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            merged = tf.get_variable(
-              f'input_{i}',
-              initializer=tf.random_normal(
-                shapes[i], mean=100, stddev=80))
-            out, out_sizes = comm.alltoallv(
-              merged, sizes[i], common_shape=[64])
-            train_ops.append(tf.group([out, out_sizes]))
-        with server.monitored_session() as sess:
-          for _ in xrange(100):
-            sess.run(train_ops)
-    finally:
-      del server
-
-  def test_multi_alltoallv(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'multi_alltoallv_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    try:
-      with tf.Graph().as_default():
-        train_ops = []
-        with tf.device(devices[0]):
-          comm0 = hb.distribute.Communicator.build(comm_id, devices)
-          ids0 = tf.constant([1, 2, 3], dtype=tf.int64)
-          sizes0 = tf.constant([1, 2], dtype=tf.int32)
-          out0, out_sizes0 = comm0.alltoallv(ids0, sizes0)
-          ids2 = tf.constant([7, 8, 9], dtype=tf.int64)
-          sizes2 = tf.constant([2, 1], dtype=tf.int32)
-          with tf.control_dependencies([out0, out_sizes0]):
-            out2, out_sizes2 = comm0.alltoallv(ids2, sizes2)
-          train_ops.append([out0, out_sizes0, out2, out_sizes2])
-        with tf.device(devices[1]):
-          comm1 = hb.distribute.Communicator.build(comm_id, devices)
-          ids1 = tf.constant([4, 5, 6], dtype=tf.int64)
-          sizes1 = tf.constant([2, 1], dtype=tf.int32)
-          out1, out_sizes1 = comm1.alltoallv(ids1, sizes1)
-          ids3 = tf.constant([3, 9, 6], dtype=tf.int64)
-          sizes3 = tf.constant([1, 2], dtype=tf.int32)
-          with tf.control_dependencies([out1, out_sizes1]):
-            out3, out_sizes3 = comm1.alltoallv(ids3, sizes3)
-          train_ops.append([out1, out_sizes1, out3, out_sizes3])
-        with server.monitored_session() as sess:
-          d0result, d1result = sess.run(train_ops)
-          np.testing.assert_equal(d0result[0], [1, 4, 5])
-          np.testing.assert_equal(d0result[1], [1, 2])
-          np.testing.assert_equal(d0result[2], [7, 8, 3])
-          np.testing.assert_equal(d0result[3], [2, 1])
-          np.testing.assert_equal(d1result[0], [2, 3, 6])
-          np.testing.assert_equal(d1result[1], [2, 1])
-          np.testing.assert_equal(d1result[2], [9, 9, 6])
-          np.testing.assert_equal(d1result[3], [1, 2])
-    finally:
-      del server
-
-  def test_mutli_alltoallv_multi_steps(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'multi_alltoallv_multistep_test'
-    num_comms = 2
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    train_ops = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          prev_op = None
-          with tf.device(d):
-            for c in xrange(num_comms):
-              shapes = [[1023 + c + 526 + c, 64], [400 + c + 99 + c, 64]]
-              sizes = [[1023 + c, 526 + c], [400 + c, 99 + c]]
-              merged = tf.get_variable(
-                f'comm_{c}/input_{i}',
-                initializer=tf.random_normal(
-                  shapes[i], mean=100, stddev=80))
-              comm = hb.distribute.Communicator.build(
-                f'{comm_id}_{c}', devices)
-              with tf.control_dependencies(
-                  None if prev_op is None else [prev_op]):
-                out, out_sizes = comm.alltoallv(
-                  merged, sizes[i], common_shape=[64])
-              train_op = tf.group([out, out_sizes])
-              prev_op = train_op
-              train_ops.append(train_op)
-        with server.monitored_session() as sess:
-          for _ in xrange(100):
-            sess.run(train_ops)
-    finally:
-      del server
-
-  def test_reuse_comm_mutli_alltoallv_multi_steps(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'reuse_comm_mutli_alltoallv_multi_steps_test'
-    num_ops = 2
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    train_ops = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            for c in xrange(num_ops):
-              shapes = [[1023 + c + 526 + c, 64], [400 + c + 99 + c, 64]]
-              sizes = [[1023 + c, 526 + c], [400 + c, 99 + c]]
-              merged = tf.get_variable(
-                f'comm_{c}/input_{i}',
-                initializer=tf.random_normal(
-                  shapes[i], mean=100, stddev=80))
-              if c == 0:
-                out, out_sizes = comm.alltoallv(
-                  merged, sizes[i], common_shape=[64])
-                train_ops.append(tf.group([out, out_sizes]))
-              else:
-                with tf.control_dependencies([train_ops[-1]]):
-                  out, out_sizes = comm.alltoallv(
-                    merged, sizes[i], common_shape=[64])
-                  train_ops.append(tf.group([out, out_sizes]))
-        with server.monitored_session() as sess:
-          for _ in xrange(100):
-            sess.run(train_ops)
-    finally:
-      del server
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/alltoallw_n_test.py b/hybridbackend/tensorflow/distribute/tests/alltoallw_n_test.py
deleted file mode 100644
index a2ec2625..00000000
--- a/hybridbackend/tensorflow/distribute/tests/alltoallw_n_test.py
+++ /dev/null
@@ -1,237 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for coalesced Alltoallw.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class AlltoallwTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_alltoallw=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def test_alltoallw_n(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_n_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    shapes_list = [
-      [[[1023, 4, 3], [526, 4, 3]], [[400, 4, 3], [99, 4, 3]]],
-      [[[34, 2, 3], [22, 2, 3]], [[11, 2, 3], [44, 2, 3]]],
-      [[[520, 5], [1314, 5]], [[888, 5], [6, 5]]]]
-    num_columns = len(shapes_list)
-    expected = []
-    actual = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(
-              comm_id, devices, impl=hb.distribute.NcclCommunicator)
-            inputs_list = []
-            for k in xrange(num_columns):
-              inputs_list.append([
-                tf.get_variable(
-                  f'from_{i}/to_{j}/var_{k}',
-                  initializer=tf.random_normal(
-                    shapes_list[k][i][j], mean=100, stddev=80))
-                for j, _ in enumerate(devices)])
-            transposed = comm.alltoallw_n(
-              inputs_list, common_shapes=[[4, 3], [2, 3], [5]])
-            actual.append(transposed)
-            expected.append(inputs_list)
-        expected = zip(*expected)
-        expected = list(zip(*[map(list, zip(*p)) for p in expected]))
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              for actj, expj in zip(acti, expi):
-                np.testing.assert_allclose(expj, actj, rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoallw_n_fp16(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    os.environ['HB_COMM_WIRE_DTYPE_FLOAT'] = 'float16'
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_n_fp16_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    shapes_list = [
-      [[[1023, 4, 3], [526, 4, 3]], [[400, 4, 3], [99, 4, 3]]],
-      [[[34, 2, 3], [22, 2, 3]], [[11, 2, 3], [44, 2, 3]]],
-      [[[520, 5], [1314, 5]], [[888, 5], [6, 5]]]]
-    num_columns = len(shapes_list)
-    expected = []
-    actual = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(
-              comm_id, devices, impl=hb.distribute.NcclCommunicator)
-            inputs_list = []
-            for k in xrange(num_columns):
-              inputs_list.append([
-                tf.get_variable(
-                  f'from_{i}/to_{j}/var_{k}',
-                  initializer=tf.random_normal(
-                    shapes_list[k][i][j], mean=100, stddev=80))
-                for j, _ in enumerate(devices)])
-            transposed = comm.alltoallw_n(
-              inputs_list, common_shapes=[[4, 3], [2, 3], [5]])
-            actual.append(transposed)
-            expected.append(inputs_list)
-        expected = zip(*expected)
-        expected = list(zip(*[map(list, zip(*p)) for p in expected]))
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              for actj, expj in zip(acti, expi):
-                np.testing.assert_allclose(expj, actj, rtol=1e-2)
-    finally:
-      del server
-    del os.environ['HB_COMM_WIRE_DTYPE_FLOAT']
-
-  def test_alltoallw_n_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_n_grad_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [[1.0, 2.4], [9.6, 8.8]]
-    shapes = [[[6, 2, 3], [8, 2, 3]], [[10, 2, 3], [12, 2, 3]]]
-    lrs = [1.0, 3.0]
-    grad_ys = [2.0]
-    all_ys = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(
-              comm_id, devices, impl=hb.distribute.NcclCommunicator)
-            inputs = [
-              tf.constant(values[i][j], shape=shapes[i][j])
-              for j, _ in enumerate(devices)]
-            all_inputs.append(inputs)
-            outputs = comm.alltoallw_n([inputs], common_shapes=[[2, 3]])
-            all_ys.append(
-              lrs[i] * tf.reduce_sum(
-                tf.concat(outputs[0], axis=0)))
-        with tf.device(devices[0]):
-          xs = [y for x in all_inputs for y in x]
-          ys = [tf.add_n(all_ys)]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-          baseline_all_outputs = list(map(list, zip(*all_inputs)))
-          baseline_all_ys = [
-            tf.reduce_sum(
-              tf.concat(baseline_all_outputs[i], axis=0)) * lrs[i]
-            for i, _ in enumerate(devices)]
-          baseline_ys = [tf.add_n(baseline_all_ys)]
-          expected = tf.gradients(
-            baseline_ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              for actj, expj in zip(acti, expi):
-                np.testing.assert_allclose(expj, actj, rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoallw_n_grad_fp16(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    os.environ['HB_COMM_WIRE_DTYPE_FLOAT'] = 'float16'
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_n_grad_fp16_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [[1.0, 2.4], [9.6, 8.8]]
-    shapes = [[[6, 2, 3], [8, 2, 3]], [[10, 2, 3], [12, 2, 3]]]
-    lrs = [1.0, 3.0]
-    grad_ys = [2.0]
-    all_ys = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(
-              comm_id, devices, impl=hb.distribute.NcclCommunicator)
-            inputs = [
-              tf.constant(values[i][j], shape=shapes[i][j])
-              for j, _ in enumerate(devices)]
-            all_inputs.append(inputs)
-            outputs = comm.alltoallw_n([inputs], common_shapes=[[2, 3]])
-            all_ys.append(
-              lrs[i] * tf.reduce_sum(
-                tf.concat(outputs[0], axis=0)))
-        with tf.device(devices[0]):
-          xs = [y for x in all_inputs for y in x]
-          ys = [tf.add_n(all_ys)]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-          baseline_all_outputs = list(map(list, zip(*all_inputs)))
-          baseline_all_ys = [
-            tf.reduce_sum(
-              tf.concat(baseline_all_outputs[i], axis=0)) * lrs[i]
-            for i, _ in enumerate(devices)]
-          baseline_ys = [tf.add_n(baseline_all_ys)]
-          expected = tf.gradients(
-            baseline_ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              for actj, expj in zip(acti, expi):
-                np.testing.assert_allclose(expj, actj, rtol=1e-2)
-    finally:
-      del server
-    del os.environ['HB_COMM_WIRE_DTYPE_FLOAT']
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/alltoallw_test.py b/hybridbackend/tensorflow/distribute/tests/alltoallw_test.py
deleted file mode 100644
index 644516ab..00000000
--- a/hybridbackend/tensorflow/distribute/tests/alltoallw_test.py
+++ /dev/null
@@ -1,276 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for Alltoallw.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class AlltoallwTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_alltoallw=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def test_alltoallw(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    shapes = [[[1023, 4, 16], [526, 4, 16]], [[400, 4, 16], [99, 4, 16]]]
-    expected = []
-    actual = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            inputs = [
-              tf.get_variable(
-                f'input_{i}/from_{j}',
-                initializer=tf.random_normal(
-                  shapes[i][j], mean=100, stddev=80))
-              for j, _ in enumerate(devices)]
-            transposed = comm.alltoallw(inputs, common_shape=[4, 16])
-            actual.append(transposed)
-            expected.append(inputs)
-        expected = list(map(list, zip(*expected)))  # transpose the inputs.
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              np.testing.assert_allclose(expi, acti, rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoallw_fp16(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    os.environ['HB_COMM_WIRE_DTYPE_FLOAT'] = 'float16'
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_fp16_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    shapes = [[[1023, 4, 16], [526, 4, 16]], [[400, 4, 16], [99, 4, 16]]]
-    expected = []
-    actual = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(
-              comm_id, devices, impl=hb.distribute.NcclCommunicator)
-            inputs = [
-              tf.get_variable(
-                f'input_{i}/from_{j}',
-                initializer=tf.random_normal(
-                  shapes[i][j], mean=100, stddev=80))
-              for j, _ in enumerate(devices)]
-            transposed = comm.alltoallw(inputs, common_shape=[4, 16])
-            actual.append(transposed)
-            expected.append(inputs)
-        expected = list(map(list, zip(*expected)))  # transpose the inputs.
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              np.testing.assert_allclose(expi, acti, rtol=1e-2)
-    finally:
-      del server
-    del os.environ['HB_COMM_WIRE_DTYPE_FLOAT']
-
-  def test_alltoallw_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_grad_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [[1.0, 2.4], [9.6, 8.8]]
-    shapes = [[[6, 2, 3], [8, 2, 3]], [[10, 2, 3], [12, 2, 3]]]
-    lrs = [1.0, 3.0]
-    grad_ys = [2.0]
-    all_ys = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            inputs = [
-              tf.constant(values[i][j], shape=shapes[i][j])
-              for j, _ in enumerate(devices)]
-            all_inputs.append(inputs)
-            outputs = comm.alltoallw(inputs, common_shape=[2, 3])
-            all_ys.append(
-              tf.reduce_sum(tf.concat(outputs, axis=0)) * lrs[i])
-        with tf.device(devices[0]):
-          xs = [y for x in all_inputs for y in x]
-          ys = [tf.add_n(all_ys)]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-          baseline_all_outputs = list(map(list, zip(*all_inputs)))
-          baseline_all_ys = [
-            tf.reduce_sum(
-              tf.concat(baseline_all_outputs[i], axis=0)) * lrs[i]
-            for i, _ in enumerate(devices)]
-          baseline_ys = [tf.add_n(baseline_all_ys)]
-          expected = tf.gradients(
-            baseline_ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              np.testing.assert_allclose(expi, acti, rtol=1e-6)
-    finally:
-      del server
-
-  def test_alltoallw_grad_fp16(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    os.environ['HB_COMM_WIRE_DTYPE_FLOAT'] = 'float16'
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_grad_fp16_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    values = [[1.0, 2.4], [9.6, 8.8]]
-    shapes = [[[6, 2, 3], [8, 2, 3]], [[10, 2, 3], [12, 2, 3]]]
-    lrs = [1.0, 3.0]
-    grad_ys = [2.0]
-    all_ys = []
-    all_inputs = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(
-              comm_id, devices, impl=hb.distribute.NcclCommunicator)
-            inputs = [
-              tf.constant(values[i][j], shape=shapes[i][j])
-              for j, _ in enumerate(devices)]
-            all_inputs.append(inputs)
-            outputs = comm.alltoallw(inputs, common_shape=[2, 3])
-            all_ys.append(
-              tf.reduce_sum(tf.concat(outputs, axis=0)) * lrs[i])
-        with tf.device(devices[0]):
-          xs = [y for x in all_inputs for y in x]
-          ys = [tf.add_n(all_ys)]
-          actual = tf.gradients(
-            ys, xs, grad_ys, colocate_gradients_with_ops=True)
-          baseline_all_outputs = list(map(list, zip(*all_inputs)))
-          baseline_all_ys = [
-            tf.reduce_sum(
-              tf.concat(baseline_all_outputs[i], axis=0)) * lrs[i]
-            for i, _ in enumerate(devices)]
-          baseline_ys = [tf.add_n(baseline_all_ys)]
-          expected = tf.gradients(
-            baseline_ys, xs, grad_ys, colocate_gradients_with_ops=True)
-        with server.monitored_session() as sess:
-          results = sess.run({'actual': actual, 'expected': expected})
-          for act, exp in zip(results['actual'], results['expected']):
-            for acti, expi in zip(act, exp):
-              np.testing.assert_allclose(expi, acti, rtol=1e-2)
-    finally:
-      del server
-    del os.environ['HB_COMM_WIRE_DTYPE_FLOAT']
-
-  def test_alltoallw_multi_steps(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'alltoallw_multistep_test'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    shapes = [[[10230, 64], [5260, 64]], [[4000, 64], [990, 64]]]
-    train_ops = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            comm = hb.distribute.Communicator.build(comm_id, devices)
-            inputs = [
-              tf.get_variable(
-                f'input_{i}/from_{j}',
-                initializer=tf.random_normal(
-                  shapes[i][j], mean=100, stddev=80))
-              for j, _ in enumerate(devices)]
-            transposed = comm.alltoallw(inputs, common_shape=[64])
-            train_op = [
-              tf.matmul(m, m, transpose_a=True) for m in transposed]
-            train_ops.append(train_op)
-        with server.monitored_session() as sess:
-          for _ in xrange(100):
-            sess.run(train_ops)
-    finally:
-      del server
-
-  def test_mutli_alltoallw_multi_steps(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    comm_id = 'multi_alltoallw_multistep_test'
-    num_comms = 2
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    train_ops = []
-    try:
-      with tf.Graph().as_default():
-        for i, d in enumerate(devices):
-          with tf.device(d):
-            for c in xrange(num_comms):
-              shapes = [
-                [[1023 + c, 64], [526 + c, 64]],
-                [[400 + c, 64], [99 + c, 64]]]
-              inputs = [
-                tf.get_variable(
-                  f'comm_{c}/input_{i}/from_{j}',
-                  initializer=tf.random_normal(
-                    shapes[i][j], mean=100, stddev=80))
-                for j, _ in enumerate(devices)]
-              comm = hb.distribute.Communicator.build(
-                f'{comm_id}_{c}', devices)
-              transposed = comm.alltoallw(inputs, common_shape=[64])
-              train_op = [
-                tf.matmul(m, m, transpose_a=True) for m in transposed]
-              train_ops.append(train_op)
-        with server.monitored_session() as sess:
-          for _ in xrange(100):
-            sess.run(train_ops)
-    finally:
-      del server
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/broadcast_in_graph_test.py b/hybridbackend/tensorflow/distribute/tests/broadcast_in_graph_test.py
deleted file mode 100644
index 40458f08..00000000
--- a/hybridbackend/tensorflow/distribute/tests/broadcast_in_graph_test.py
+++ /dev/null
@@ -1,73 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for Broadcast in one graph.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class BroadcastInGraphTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_broadcast=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def test_broadcast_in_graph(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    a = 13
-    b = 22
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        root = tf.constant(a)
-        recv0 = comm0.broadcast(root, root_rank=0)
-      with tf.device('/gpu:1'):
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        noop = tf.constant(b)
-        recv1 = comm1.broadcast(noop, root_rank=0)
-      with tf.train.MonitoredTrainingSession('') as sess:
-        s0, s1 = sess.run([recv0, recv1])
-        np.testing.assert_allclose(s0, a, rtol=1e-6)
-        np.testing.assert_allclose(s1, a, rtol=1e-6)
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/broadcast_test.py b/hybridbackend/tensorflow/distribute/tests/broadcast_test.py
index 28e17d17..169cc9e5 100644
--- a/hybridbackend/tensorflow/distribute/tests/broadcast_test.py
+++ b/hybridbackend/tensorflow/distribute/tests/broadcast_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # =============================================================================
 
-r'''Tests for Broadcast.
+r'''Tests for broadcast collective communication.
 '''
 
 from __future__ import absolute_import
@@ -27,25 +27,24 @@
 
 import hybridbackend.common.test as hbtest
 
+# pylint: disable=missing-docstring,import-outside-toplevel
+
 
 def _test_broadcast(rank, a, b):
   r'''Test Broadcast.
   '''
-  # pylint: disable=import-outside-toplevel
   import tensorflow as tf
 
   import hybridbackend.tensorflow as hb
 
   with tf.Graph().as_default():
     with hb.scope():
-      comm = hb.distribute.Communicator.build('comm0', hb.context.devices)
-      root = tf.constant(a) if rank == 0 else tf.constant(b)
-      recv = comm.broadcast(root, root_rank=0)
+      data = tf.constant(a) if rank == 0 else tf.constant(b)
+      recv = hb.distribute.broadcast(data, root_rank=0)
       with tf.train.MonitoredTrainingSession('') as sess:
         return sess.run(recv)
 
 
-# pylint: disable=missing-docstring
 @unittest.skipUnless(
   os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
 @unittest.skipUnless(
@@ -55,7 +54,10 @@ def setUp(self):  # pylint: disable=invalid-name
     os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
     os.environ['NCCL_DEBUG'] = 'INFO'
     os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_broadcast=1'
+    os.environ['TF_CPP_VMODULE'] = (
+      'nccl_comm=1,'
+      'nccl_create=1,'
+      'nccl_broadcast=1')
 
   def tearDown(self):  # pylint: disable=invalid-name
     del os.environ['TF_CPP_VMODULE']
diff --git a/hybridbackend/tensorflow/ops/tests/floormod_shuffle_test.py b/hybridbackend/tensorflow/distribute/tests/partition_test.py
similarity index 82%
rename from hybridbackend/tensorflow/ops/tests/floormod_shuffle_test.py
rename to hybridbackend/tensorflow/distribute/tests/partition_test.py
index 7474100c..9beec7e2 100644
--- a/hybridbackend/tensorflow/ops/tests/floormod_shuffle_test.py
+++ b/hybridbackend/tensorflow/distribute/tests/partition_test.py
@@ -33,7 +33,7 @@
 # pylint: disable=missing-docstring
 @unittest.skipUnless(
   os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-class FloorModShuffleTest(unittest.TestCase):
+class PartitionTest(unittest.TestCase):
   def setUp(self):  # pylint: disable=invalid-name
     os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 
@@ -48,8 +48,8 @@ def test_unfused(self):
 
     with tf.Graph().as_default(), hb.scope():
       xi = tf.constant(x)
-      y_list, idx_list = hb.floormod_partition(xi, num_partitions)
-      y, ysizes, idx = hb.floormod_shuffle(xi, num_partitions)
+      y_list, idx_list = hb.distribute.partition(xi, num_partitions)
+      y, ysizes, idx = hb.distribute.partition_by_modulo(xi, num_partitions)
       with tf.train.MonitoredTrainingSession('') as sess:
         tf_y_list, tf_idx_list, tf_y, tf_ysizes, tf_idx = sess.run(
           [y_list, idx_list, y, ysizes, idx])
@@ -63,9 +63,6 @@ def test_unfused(self):
       np.take(
         x,
         np.concatenate(tf_idx_list, axis=0)))
-    np.testing.assert_equal(
-      np.concatenate(tf_y_list, axis=0),
-      tf_y)
 
   def test_empty_input(self):
     x = np.array([], np.int64)
@@ -74,7 +71,7 @@ def test_empty_input(self):
     with tf.Graph().as_default(), hb.scope():
       with tf.device('/gpu:0'):
         xi = tf.constant(x)
-        y, ysizes, idx = hb.floormod_shuffle(xi, num_partitions)
+        y, ysizes, idx = hb.distribute.partition_by_modulo(xi, num_partitions)
       with tf.train.MonitoredTrainingSession('') as sess:
         tf_y, tf_ysizes, tf_idx = sess.run([y, ysizes, idx])
 
@@ -98,8 +95,16 @@ def test_fused(self):
     with tf.Graph().as_default(), hb.scope():
       with tf.device('/gpu:0'):
         xt = [tf.constant(t) for t in x]
-        y, ysizes, idx = hb.floormod_shuffle_n(
-          xt, num_partitions)
+        y = []
+        ysizes = []
+        idx = []
+        for xt_item in xt:
+          y_out, ysizes_out, idx_out = hb.distribute.partition_by_modulo(
+            xt_item, num_partitions)
+          y.append(y_out)
+          ysizes.append(ysizes_out)
+          idx.append(idx_out)
+
         with tf.train.MonitoredTrainingSession('') as sess:
           result = sess.run({'y': y, 'ysizes': ysizes, 'idx': idx})
 
@@ -116,8 +121,15 @@ def test_empty_input_fused(self):
     with tf.Graph().as_default(), hb.scope():
       with tf.device('/gpu:0'):
         xt = [tf.constant(t) for t in x]
-        y, ysizes, idx = hb.floormod_shuffle_n(
-          xt, num_partitions)
+        y = []
+        ysizes = []
+        idx = []
+        for xt_item in xt:
+          y_out, ysizes_out, idx_out = hb.distribute.partition_by_modulo(
+            xt_item, num_partitions)
+          y.append(y_out)
+          ysizes.append(ysizes_out)
+          idx.append(idx_out)
       with tf.train.MonitoredTrainingSession('') as sess:
         result = sess.run({'y': y, 'ysizes': ysizes, 'idx': idx})
 
diff --git a/hybridbackend/tensorflow/distribute/tests/reduce_scatter_test.py b/hybridbackend/tensorflow/distribute/tests/reduce_scatter_test.py
deleted file mode 100644
index d4806e19..00000000
--- a/hybridbackend/tensorflow/distribute/tests/reduce_scatter_test.py
+++ /dev/null
@@ -1,167 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for ReduceScatter.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class ReduceScatterTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_reduce_scatter=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def _test_reduce_scatter(self, devices, shapes, islegal=True):
-    hb.context.options.update(comm_pubsub_device='')
-
-    num_comms = len(shapes)
-    num_devices = len(devices)
-    all_inputs = []
-    all_sums = []
-
-    with tf.Graph().as_default(), hb.scope():
-      for i in xrange(num_comms):
-        comm_inputs = []
-        comm_sums = []
-        shared_name = f'comm_{i}'
-        for d in xrange(num_devices):
-          with tf.device(devices[d]):
-            comm_inputs.append(
-              tf.get_variable(
-                f'input{i}/part_{d}',
-                initializer=tf.random_normal(
-                  shapes[i], mean=100, stddev=80)))
-            comm = hb.distribute.Communicator.build(shared_name, devices)
-            comm_sums.append(comm.reduce_scatter(comm_inputs[d]))
-        all_inputs.append(comm_inputs)
-        all_sums.append(comm_sums)
-      baselines = (
-        [tf.split(tf.add_n(all_inputs[i]), num_devices, axis=0)
-         for i in xrange(num_comms)]
-        if islegal else [])
-      with tf.train.MonitoredTrainingSession('') as sess:
-        expects = sess.run(baselines)
-        actuals = sess.run(all_sums)
-      for i in xrange(num_comms):
-        for d in xrange(num_devices):
-          np.testing.assert_allclose(actuals[i][d], expects[i][d], atol=1e-4)
-
-  def test_one_device_one_tensor_reduce_scatter(self):
-    self._test_reduce_scatter(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0]],
-      [[16, 4]])
-
-  def test_two_devices_scalar_tensors_reduce_scatter(self):
-    with self.assertRaises(ValueError):
-      self._test_reduce_scatter(
-        [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0, 1]],
-        [[]])
-      self.fail('Should throw ValueError here.')
-
-  def test_two_devices_four_tensors_reduce_scatter(self):
-    self._test_reduce_scatter(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0, 1]],
-      [[18, 41, 37], [16, 121], [2], [16, 6]])
-
-  def test_two_devices_six_tensors_reduce_scatter(self):
-    self._test_reduce_scatter(
-      [f'/job:localhost/replica:0/task:0/device:GPU:{d}'
-       for d in [0, 1]],
-      [[16, 4], [102, 100], [4, 5, 8], [12], [6, 2, 2, 3], [1048]])
-
-  def test_two_devices_irregular_reduce_scatter(self):
-    with self.assertRaises(tf.errors.InvalidArgumentError):
-      self._test_reduce_scatter(
-        [f'/job:localhost/replica:0/task:0/device:GPU:{d}' for d in [0, 1]],
-        [[17, 41, 37]], False)
-      self.fail('Should throw InvalidArgumentError here.')
-
-  def test_reduce_scatter_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[4, 8])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.reduce_scatter(a * 0.5)
-        loss0 = tf.reduce_mean(recv0) * 20.0
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[4, 8])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        recv1 = comm1.reduce_scatter(b * 0.75)
-        loss1 = tf.reduce_mean(recv1) * 10.0
-      loss = loss0 * loss1 + loss1 * 2.0
-      grad0, grad1 = tf.gradients([loss], [a, b], [2.0])
-      with tf.train.MonitoredTrainingSession('') as sess:
-        g0, g1 = sess.run([grad0, grad1])
-        np.testing.assert_allclose(
-          g0,
-          [[25.0] * 8, [25.0] * 8, [26.25] * 8, [26.25] * 8],
-          rtol=1e-6)
-        np.testing.assert_allclose(
-          g1,
-          [[37.5] * 8, [37.5] * 8, [39.375] * 8, [39.375] * 8],
-          rtol=1e-6)
-
-  def test_reduce_scatter_grad_error(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default():
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[4, 8])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.reduce_scatter(a * 0.5, reduce_op=hb.distribute.ops.MAX)
-        loss0 = tf.reduce_mean(recv0) * 20.0
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[4, 8])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        recv1 = comm1.reduce_scatter(b * 0.75, reduce_op=hb.distribute.ops.MAX)
-        loss1 = tf.reduce_mean(recv1) * 10.0
-      loss = loss0 * loss1 + loss1 * 2.0
-      with self.assertRaises(RuntimeError):
-        tf.gradients([loss], [a, b], [2.0], colocate_gradients_with_ops=True)
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/distribute/tests/reduce_test.py b/hybridbackend/tensorflow/distribute/tests/reduce_test.py
deleted file mode 100644
index 15c4fbdf..00000000
--- a/hybridbackend/tensorflow/distribute/tests/reduce_test.py
+++ /dev/null
@@ -1,209 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for Reduce.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-from six.moves import xrange  # pylint: disable=redefined-builtin
-import unittest
-
-import numpy as np
-import tensorflow as tf
-
-import hybridbackend.common.test as hbtest
-import hybridbackend.tensorflow as hb
-
-
-# pylint: disable=missing-docstring
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-class ReduceTest(unittest.TestCase):
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['NCCL_DEBUG'] = 'INFO'
-    os.environ['NCCL_DEBUG_SUBSYS'] = 'ALL'
-    os.environ['NCCL_LAUNCH_MODE'] = 'GROUP'
-    os.environ['TF_CPP_VMODULE'] = 'nccl_reduce=1'
-
-  def tearDown(self):  # pylint: disable=invalid-name
-    del os.environ['TF_CPP_VMODULE']
-    del os.environ['CUDA_VISIBLE_DEVICES']
-
-  def _simple_reduce(self, root_rank=0):
-    hb.context.options.update(comm_pubsub_device='')
-
-    a = 13
-    b = 22
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        input0 = tf.constant(a)
-        sum0 = comm0.reduce(input0, root_rank=root_rank)
-      with tf.device('/gpu:1'):
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        input1 = tf.constant(b)
-        sum1 = comm1.reduce(input1, root_rank=root_rank)
-      with tf.train.MonitoredTrainingSession('') as sess:
-        s0, s1 = sess.run([sum0, sum1])
-        if root_rank == 0:
-          np.testing.assert_allclose(s0, a + b, rtol=1e-6)
-        else:
-          np.testing.assert_allclose(s1, a + b, rtol=1e-6)
-
-  def test_simple_reduce(self):
-    self._simple_reduce(root_rank=0)
-    self._simple_reduce(root_rank=1)
-
-  def test_onedevice_reduce(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    a = 13
-    devices = ['/gpu:0']
-    shared_name = 'comm_onedevice'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        input0 = tf.constant(a)
-        sum0 = comm0.reduce(input0)
-      with tf.train.MonitoredTrainingSession('') as sess:
-        np.testing.assert_allclose(sess.run(sum0), a, rtol=1e-6)
-
-  def test_reduce_max(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    a = 13
-    b = 22
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        input0 = tf.constant(a)
-        max0 = comm0.reduce(input0, reduce_op=hb.distribute.ops.MAX)
-      with tf.device('/gpu:1'):
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        input1 = tf.constant(b)
-        max1 = comm1.reduce(input1, reduce_op=hb.distribute.ops.MAX)
-      with tf.train.MonitoredTrainingSession('') as sess:
-        s0, _ = sess.run([max0, max1])
-        np.testing.assert_allclose(s0, b, rtol=1e-6)
-
-  def test_multicomm(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    a = 13
-    b = 22
-    devices = ['/gpu:0', '/gpu:1']
-    with tf.Graph().as_default(), hb.scope():
-      reduce_ops_rank0 = []
-      reduce_ops_rank1 = []
-      for icomm in xrange(10):
-        comm_name = f'comm{icomm}'
-        with tf.device('/gpu:0'):
-          comm0 = hb.distribute.Communicator.build(comm_name, devices)
-          input0 = tf.constant(a)
-          sum0 = comm0.reduce(input0)
-        with tf.device('/gpu:1'):
-          comm1 = hb.distribute.Communicator.build(comm_name, devices)
-          input1 = tf.constant(b)
-          sum1 = comm1.reduce(input1)
-        reduce_ops_rank0.append(sum0)
-        reduce_ops_rank1.append(sum1)
-      with tf.train.MonitoredTrainingSession('') as sess:
-        results_rank0, _ = sess.run([reduce_ops_rank0, reduce_ops_rank1])
-        for r in results_rank0:
-          np.testing.assert_allclose(r, a + b, rtol=1e-6)
-
-  def test_monitored_training_session(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    a = 13
-    b = 22
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    server = hb.train.Server({'localhost': ['localhost:0']})
-    try:
-      with tf.Graph().as_default():
-        with tf.device('/gpu:0'):
-          comm0 = hb.distribute.Communicator.build(shared_name, devices)
-          input0 = tf.constant(a)
-          sum0 = comm0.reduce(input0)
-        with tf.device('/gpu:1'):
-          comm1 = hb.distribute.Communicator.build(shared_name, devices)
-          input1 = tf.constant(b)
-          sum1 = comm1.reduce(input1)
-
-        with server.monitored_session() as sess:
-          s0, _ = sess.run([sum0, sum1])
-          np.testing.assert_allclose(s0, a + b, rtol=1e-6)
-    finally:
-      del server
-
-  def test_reduce_grad(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default(), hb.scope():
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[2, 10])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.reduce(a * 0.75)
-        loss0 = tf.reduce_mean(recv0) * 20.0
-        grad0 = tf.gradients(
-          [loss0], [a], [2.0],
-          colocate_gradients_with_ops=True)
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[2, 10])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        recv1 = comm1.reduce(b)
-        loss1 = tf.reduce_mean(recv1) * 20.0
-        grad1 = tf.gradients(
-          [loss1], [b], [2.0],
-          colocate_gradients_with_ops=True)
-      with tf.train.MonitoredTrainingSession('') as sess:
-        vals = sess.run({'g0': grad0, 'g1': grad1})
-        np.testing.assert_allclose(vals['g0'][0], [[1.5] * 10] * 2, rtol=1e-6)
-
-  def test_reduce_grad_error(self):
-    hb.context.options.update(comm_pubsub_device='')
-
-    devices = ['/gpu:0', '/gpu:1']
-    shared_name = 'comm'
-    with tf.Graph().as_default():
-      with tf.device('/gpu:0'):
-        a = tf.constant(1.0, shape=[2, 10])
-        comm0 = hb.distribute.Communicator.build(shared_name, devices)
-        recv0 = comm0.reduce(a * 0.75, reduce_op=hb.distribute.ops.MAX)
-        loss = tf.reduce_mean(recv0) * 20.0
-      with tf.device('/gpu:1'):
-        b = tf.constant(2.0, shape=[2, 10])
-        comm1 = hb.distribute.Communicator.build(shared_name, devices)
-        _ = comm1.reduce(b, reduce_op=hb.distribute.ops.MAX)
-      with self.assertRaises(RuntimeError):
-        tf.gradients([loss], [a], [2.0])
-
-
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/feature_column/tests/dense_features_test.py b/hybridbackend/tensorflow/feature_column/tests/dense_features_test.py
deleted file mode 100644
index 35753a6f..00000000
--- a/hybridbackend/tensorflow/feature_column/tests/dense_features_test.py
+++ /dev/null
@@ -1,673 +0,0 @@
-# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# =============================================================================
-
-r'''Tests for embedding columns.
-'''
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-import unittest
-
-import numpy as np
-
-import hybridbackend.common.test as hbtest
-
-
-# pylint: disable=missing-docstring
-# pylint: disable=import-outside-toplevel
-def _test_get_dense_tensor(_):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  # Inputs.
-  vocabulary_size = 3
-  sparse_input = tf.SparseTensorValue(
-    indices=((0, 0), (1, 0), (1, 4), (3, 0)),
-    values=(2, 0, 1, 1),
-    dense_shape=(4, 5))
-
-  # Embedding variable.
-  embedding_dimension = 2
-  embedding_values = (
-    (1., 2.),  # id 0
-    (3., 5.),  # id 1
-    (7., 11.)  # id 2
-  )
-
-  def _initializer(shape, dtype, partition_info):
-    np.testing.assert_equal((vocabulary_size, embedding_dimension), shape)
-    np.testing.assert_equal(tf.float32, dtype)
-    np.testing.assert_equal(partition_info, None)
-    return embedding_values
-
-  with tf.Graph().as_default():
-    with hb.scope():
-      # Build columns.
-      categorical_column = tf.feature_column.categorical_column_with_identity(
-        key='aaa', num_buckets=vocabulary_size)
-      emb_col = tf.feature_column.embedding_column(
-        categorical_column,
-        embedding_dimension,
-        initializer=_initializer,
-        combiner='mean')
-      # Provide sparse input and get dense result.
-      embedding_lookup = hb.keras.layers.DenseFeatures(
-        [emb_col])({'aaa': sparse_input})
-      with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(embedding_lookup)
-
-
-def _test_get_dense_tensor_sharded(rank):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  # Inputs.
-  vocabulary_size = 3
-  sparse_input = tf.SparseTensorValue(
-    indices=((0, 0), (1, 0), (1, 4), (3, 0)),
-    values=(2, 0, 1, 1),
-    dense_shape=(4, 5))
-
-  # Embedding variable.
-  embedding_dimension = 2
-  embedding_values_0 = (
-    (1., 2.),  # id 0
-    (7., 11.),  # id 2
-  )
-  embedding_values_1 = (
-    (3., 5.),  # id 1
-  )
-
-  def _initializer_0(shape, dtype, partition_info):
-    np.testing.assert_equal((2, embedding_dimension), shape)
-    np.testing.assert_equal(tf.float32, dtype)
-    np.testing.assert_equal(partition_info, None)
-    return embedding_values_0
-
-  def _initializer_1(shape, dtype, partition_info):
-    np.testing.assert_equal((1, embedding_dimension), shape)
-    np.testing.assert_equal(tf.float32, dtype)
-    np.testing.assert_equal(partition_info, None)
-    return embedding_values_1
-
-  with tf.Graph().as_default():
-    with hb.scope():
-      categorical_column = tf.feature_column.categorical_column_with_identity(
-        key='aaa', num_buckets=vocabulary_size)
-      emb_col = tf.feature_column.embedding_column(
-        categorical_column,
-        embedding_dimension,
-        initializer=_initializer_0 if rank == 0 else _initializer_1,
-        combiner='mean')
-      # Provide sparse input and get dense result.
-      embedding_lookup = hb.keras.layers.DenseFeatures(
-        [emb_col])({'aaa': sparse_input})
-      with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(embedding_lookup)
-
-
-def _test_get_dense_tensor_with_varscope(rank):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  # Inputs.
-  vocabulary_size = 3
-  sparse_input = tf.SparseTensorValue(
-    indices=((0, 0), (1, 0), (1, 4), (3, 0)),
-    values=(2, 0, 1, 1),
-    dense_shape=(4, 5))
-
-  # Embedding variable.
-  embedding_dimension = 2
-  embedding_values_0 = (
-    (1., 2.),  # id 0
-    (7., 11.),  # id 2
-  )
-  embedding_values_1 = (
-    (3., 5.),  # id 1
-  )
-
-  def _initializer_0(shape, dtype, partition_info):
-    np.testing.assert_equal((1, embedding_dimension), shape)
-    np.testing.assert_equal(tf.float32, dtype)
-    return [embedding_values_0[partition_info.var_offset[0]]]
-
-  def _initializer_1(shape, dtype, partition_info):
-    del partition_info
-    np.testing.assert_equal((1, embedding_dimension), shape)
-    np.testing.assert_equal(tf.float32, dtype)
-    return embedding_values_1
-
-  with tf.Graph().as_default():
-    with hb.scope():
-      partitioner = tf.min_max_variable_partitioner(
-        max_partitions=2, min_slice_size=4)
-      with tf.variable_scope('test', partitioner=partitioner):
-        categorical_column = tf.feature_column.categorical_column_with_identity(
-          key='aaa', num_buckets=vocabulary_size)
-        emb_col = tf.feature_column.embedding_column(
-          categorical_column,
-          embedding_dimension,
-          initializer=_initializer_0 if rank == 0 else _initializer_1,
-          combiner='mean')
-        # Provide sparse input and get dense result.
-        embedding_lookup = hb.keras.layers.DenseFeatures(
-          [emb_col])({'aaa': sparse_input})
-        with tf.train.MonitoredTrainingSession('') as sess:
-          return sess.run(embedding_lookup)
-
-
-def _test_embedding_column_with_optimizer(_, lr):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(seed=42):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad0', num_buckets=10, default_value=0),
-          dimension=20,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad1', num_buckets=10, default_value=0),
-          dimension=30,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad2', num_buckets=10, default_value=0),
-          dimension=40,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='user0', num_buckets=10, default_value=0),
-          dimension=20,
-          initializer=tf.constant_initializer(0.5)),
-      ]
-      features = {
-        'ad0': tf.constant([0, 1, 3, 2]),
-        'ad1': tf.constant([1, 5, 3, 4]),
-        'ad2': tf.constant([5, 2, 7, 4]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      out_emb = hb.keras.layers.DenseFeatures(
-        columns, num_groups=None)(features)
-      loss = tf.reduce_mean(out_emb)
-      opt = tf.train.AdamOptimizer(lr)
-      step = tf.train.get_or_create_global_step()
-      train_op = opt.minimize(loss, global_step=step)
-
-      final_loss = None
-      with tf.train.MonitoredTrainingSession(
-          '',
-          hooks=[
-            tf.train.StopAtStepHook(last_step=100),
-            tf.train.NanTensorHook(loss),
-            tf.train.LoggingTensorHook(
-              tensors={'loss': loss, 'step': step},
-              every_n_iter=20)]) as sess:
-        while not sess.should_stop():
-          final_loss = sess.run(loss)
-          sess.run(train_op)
-      return final_loss
-
-
-def _test_get_dense_tensor_disable_concat(_):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(emb_enable_concat=False):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad0', num_buckets=10, default_value=0),
-          dimension=20,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='user0', num_buckets=10, default_value=0),
-          dimension=20,
-          initializer=tf.constant_initializer(0.5)),
-      ]
-      features = {
-        'ad0': tf.sparse.SparseTensor(
-          values=[0, 1, 3, 2],
-          indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
-          dense_shape=[2, 2]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      embs = hb.keras.layers.DenseFeatures(columns)(features)
-
-      with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(embs)
-
-
-def _test_embedding_column_with_coalescing(_, lr):
-  os.environ['HYBRIDBACKEND_DEFAULT_COMM'] = 'NCCL'
-
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(seed=42):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad0', num_buckets=10, default_value=0),
-          dimension=20,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad1', num_buckets=10, default_value=0),
-          dimension=30,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad2', num_buckets=10, default_value=0),
-          dimension=40,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='user0', num_buckets=10, default_value=0),
-          dimension=20,
-          initializer=tf.constant_initializer(0.5)),
-      ]
-      features = {
-        'ad0': tf.constant([0, 1, 3, 2]),
-        'ad1': tf.constant([1, 5, 3, 4]),
-        'ad2': tf.constant([5, 2, 7, 4]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      out_emb = hb.keras.layers.DenseFeatures(columns, num_groups=2)(features)
-      loss = tf.reduce_mean(out_emb)
-      opt = tf.train.AdamOptimizer(lr)
-      step = tf.train.get_or_create_global_step()
-      train_op = opt.minimize(loss, global_step=step)
-
-      final_loss = None
-      with tf.train.MonitoredTrainingSession(
-          '',
-          hooks=[
-            tf.train.StopAtStepHook(last_step=100),
-            tf.train.NanTensorHook(loss),
-            tf.train.LoggingTensorHook(
-              tensors={'loss': loss, 'step': step},
-              every_n_iter=20)]) as sess:
-        while not sess.should_stop():
-          final_loss = sess.run(loss)
-          sess.run(train_op)
-      return final_loss
-
-
-def _test_embedding_column_with_function(_, lr):
-  os.environ['HYBRIDBACKEND_DEFAULT_COMM'] = 'NCCL'
-
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  def train_fn():
-    columns = [
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='ad0', num_buckets=10, default_value=0),
-        dimension=20,
-        initializer=tf.constant_initializer(0.5)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='ad1', num_buckets=10, default_value=0),
-        dimension=30,
-        initializer=tf.constant_initializer(0.5)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='ad2', num_buckets=10, default_value=0),
-        dimension=40,
-        initializer=tf.constant_initializer(0.5)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='user0', num_buckets=10, default_value=0),
-        dimension=20,
-        initializer=tf.constant_initializer(0.5)),
-    ]
-    features = {
-      'ad0': tf.constant([0, 1, 3, 2]),
-      'ad1': tf.constant([1, 5, 3, 4]),
-      'ad2': tf.constant([5, 2, 7, 4]),
-      'user0': tf.constant([2, 5, 4, 7])
-    }
-    out_emb = hb.keras.layers.DenseFeatures(columns)(features)
-    loss = tf.reduce_mean(out_emb)
-    opt = tf.train.AdamOptimizer(lr)
-    step = tf.train.get_or_create_global_step()
-    return loss, opt.minimize(loss, global_step=step)
-
-  with hb.scope(seed=42, emb_num_groups=2):
-    loss, train_op = train_fn()
-    final_loss = None
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=[
-          tf.train.StopAtStepHook(last_step=100),
-          tf.train.NanTensorHook(loss),
-          tf.train.LoggingTensorHook(
-            tensors={'loss': loss},
-            every_n_iter=20)]) as sess:
-      while not sess.should_stop():
-        final_loss = sess.run(loss)
-        sess.run(train_op)
-    return final_loss
-
-
-def _test_embedding_column_with_function_unique(_, lr):
-  os.environ['HYBRIDBACKEND_DEFAULT_COMM'] = 'NCCL'
-
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  def train_fn():
-    columns = [
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='ad0', num_buckets=10, default_value=0),
-        dimension=20,
-        initializer=tf.constant_initializer(0.5)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='ad1', num_buckets=10, default_value=0),
-        dimension=30,
-        initializer=tf.constant_initializer(0.5)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='ad2', num_buckets=10, default_value=0),
-        dimension=40,
-        initializer=tf.constant_initializer(0.5)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          key='user0', num_buckets=10, default_value=0),
-        dimension=20,
-        initializer=tf.constant_initializer(0.5)),
-    ]
-    features = {
-      'ad0': tf.constant([0, 1, 3, 2]),
-      'ad1': tf.constant([1, 5, 3, 4]),
-      'ad2': tf.constant([5, 2, 7, 4]),
-      'user0': tf.constant([2, 5, 4, 7])
-    }
-    out_emb = hb.keras.layers.DenseFeatures(columns)(features)
-    loss = tf.reduce_mean(out_emb)
-    opt = tf.train.AdamOptimizer(lr)
-    step = tf.train.get_or_create_global_step()
-    return loss, opt.minimize(loss, global_step=step)
-
-  with hb.scope(seed=42, emb_num_groups=2, emb_unique={'ad0': True}):
-    loss, train_op = train_fn()
-    final_loss = None
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=[
-          tf.train.StopAtStepHook(last_step=100),
-          tf.train.NanTensorHook(loss),
-          tf.train.LoggingTensorHook(
-            tensors={'loss': loss},
-            every_n_iter=20)]) as sess:
-      while not sess.should_stop():
-        final_loss = sess.run(loss)
-        sess.run(train_op)
-    return final_loss
-
-
-def _test_get_dense_tensor_with_segment_rank(rank):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  # Inputs.
-  vocabulary_size = 3
-  sparse_input = tf.SparseTensorValue(
-    indices=((0, 1, 1), (0, 1, 2), (1, 1, 1), (1, 1, 2)),
-    values=(2, 0, 1, 1),
-    dense_shape=(4, 2, 3))
-
-  # Embedding variable.
-  embedding_dimension = 2
-  embedding_values_0 = (
-    (1., 2.),  # id 0
-    (7., 11.),  # id 2
-  )
-  embedding_values_1 = (
-    (3., 5.),  # id 1
-  )
-
-  def _initializer_0(shape, dtype, partition_info):
-    np.testing.assert_equal((2, embedding_dimension), shape)
-    np.testing.assert_equal(tf.float32, dtype)
-    np.testing.assert_equal(partition_info, None)
-    return embedding_values_0
-
-  def _initializer_1(shape, dtype, partition_info):
-    np.testing.assert_equal((1, embedding_dimension), shape)
-    np.testing.assert_equal(tf.float32, dtype)
-    np.testing.assert_equal(partition_info, None)
-    return embedding_values_1
-
-  @hb.function(emb_segment_rank={'aaa': 1})
-  def lookup_fn():
-    categorical_column = tf.feature_column.categorical_column_with_identity(
-      key='aaa', num_buckets=vocabulary_size)
-    emb_col = tf.feature_column.embedding_column(
-      categorical_column,
-      embedding_dimension,
-      initializer=_initializer_0 if rank == 0 else _initializer_1,
-      combiner='mean')
-    return hb.keras.layers.DenseFeatures([emb_col])({'aaa': sparse_input})
-
-  with tf.Graph().as_default(), hb.scope():
-    embs = lookup_fn()
-    with tf.train.MonitoredTrainingSession('') as sess:
-      return sess.run(embs)
-
-
-def _test_shared_embedding_column(_, lr):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(seed=42):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad1', num_buckets=10, default_value=0),
-          dimension=30,
-          initializer=tf.constant_initializer(0.5)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad2', num_buckets=10, default_value=0),
-          dimension=40,
-          initializer=tf.constant_initializer(0.5)),
-      ]
-      columns += tf.feature_column.shared_embedding_columns(
-        [
-          tf.feature_column.categorical_column_with_identity(
-            key='ad0', num_buckets=10, default_value=0),
-          tf.feature_column.categorical_column_with_identity(
-            key='user0', num_buckets=10, default_value=0)],
-        dimension=20,
-        initializer=tf.constant_initializer(0.5))
-      features = {
-        'ad0': tf.constant([0, 1, 3, 2]),
-        'ad1': tf.constant([1, 5, 3, 4]),
-        'ad2': tf.constant([5, 2, 7, 4]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      out_emb = hb.keras.layers.DenseFeatures(columns)(features)
-      loss = tf.reduce_mean(out_emb)
-      opt = tf.train.AdamOptimizer(lr)
-      step = tf.train.get_or_create_global_step()
-      train_op = opt.minimize(loss, global_step=step)
-
-      final_loss = None
-      with tf.train.MonitoredTrainingSession(
-          '',
-          hooks=[
-            tf.train.StopAtStepHook(last_step=100),
-            tf.train.NanTensorHook(loss),
-            tf.train.LoggingTensorHook(
-              tensors={'loss': loss, 'step': step},
-              every_n_iter=20)]) as sess:
-        while not sess.should_stop():
-          final_loss = sess.run(loss)
-          sess.run(train_op)
-      return final_loss
-
-
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_CUDA') == 'ON', 'GPU required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_NCCL') == 'ON', 'NCCL required')
-@unittest.skipUnless(
-  os.getenv('HYBRIDBACKEND_WITH_TENSORFLOW_ESTIMATOR') != 'OFF',
-  'TF Estimator required')
-class DenseFeaturesTest(unittest.TestCase):
-  '''Tests for embedding columns.
-  '''
-  def setUp(self):  # pylint: disable=invalid-name
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-    os.environ['TF_CPP_VMODULE'] = (
-      'op_optimization=2,'
-      'replacing=2,'
-      'packing=2,'
-      'pruning=2,'
-      'relocation=2')
-
-  def test_get_dense_tensor(self):
-    results = hbtest.Spawn()(_test_get_dense_tensor)
-    np.testing.assert_allclose(
-      results[0],
-      [[7., 11.],
-       [2., 3.5],
-       [0., 0.],
-       [3., 5.]],
-      rtol=1e-6)
-
-  def test_get_dense_tensor_sharded(self):
-    results = hbtest.Spawn(2)(_test_get_dense_tensor_sharded)
-    np.testing.assert_allclose(
-      results[0],
-      [[7., 11.],
-       [2., 3.5],
-       [0., 0.],
-       [3., 5.]],
-      rtol=1e-6)
-    np.testing.assert_allclose(
-      results[1],
-      [[7., 11.],
-       [2., 3.5],
-       [0., 0.],
-       [3., 5.]],
-      rtol=1e-6)
-
-  def test_get_dense_tensor_with_varscope(self):
-    results = hbtest.Spawn(2)(_test_get_dense_tensor_with_varscope)
-    np.testing.assert_allclose(
-      results[0],
-      [[7., 11.],
-       [2., 3.5],
-       [0., 0.],
-       [3., 5.]],
-      rtol=1e-6)
-    np.testing.assert_allclose(
-      results[1],
-      [[7., 11.],
-       [2., 3.5],
-       [0., 0.],
-       [3., 5.]],
-      rtol=1e-6)
-
-  def test_embedding_column_with_optimizer(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_optimizer(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.490101, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.490101, rtol=1e-6)
-
-  def test_get_dense_tensor_disable_concat(self):
-    results = hbtest.Spawn()(_test_get_dense_tensor_disable_concat)
-    np.testing.assert_equal(len(results[0]), 2)
-
-  def test_embedding_column_with_coalescing(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_coalescing(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.490101, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.490101, rtol=1e-6)
-
-  def test_embedding_column_function(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_function(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.490101, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.490101, rtol=1e-6)
-
-  def test_embedding_column_function_unique(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_function_unique(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.490101, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.490101, rtol=1e-6)
-
-  def test_get_dense_tensor_with_segment_rank(self):
-    results = hbtest.Spawn(2)(_test_get_dense_tensor_with_segment_rank)
-    np.testing.assert_allclose(
-      results[0],
-      [[0., 0.],
-       [4., 6.5],
-       [0., 0.],
-       [3., 5.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.]],
-      rtol=1e-6)
-    np.testing.assert_allclose(
-      results[1],
-      [[0., 0.],
-       [4., 6.5],
-       [0., 0.],
-       [3., 5.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.]],
-      rtol=1e-6)
-
-  def test_shared_embedding_column(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_shared_embedding_column(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.490101, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.490101, rtol=1e-6)
-
-
-# pylint: enable=missing-docstring
-if __name__ == '__main__':
-  hbtest.main(f'{__file__}.xml')
diff --git a/hybridbackend/tensorflow/plugins/deeprec/tests/ev_test.py b/hybridbackend/tensorflow/plugins/deeprec/tests/ev_test.py
index efb707e0..61c92a3e 100644
--- a/hybridbackend/tensorflow/plugins/deeprec/tests/ev_test.py
+++ b/hybridbackend/tensorflow/plugins/deeprec/tests/ev_test.py
@@ -21,467 +21,41 @@
 from __future__ import print_function
 
 import os
-import sys
 import unittest
 
-import numpy as np
-
 import hybridbackend.common.test as hbtest
 
-
 # pylint: disable=missing-docstring
 # pylint: disable=import-outside-toplevel
-def _test_get_dense_tensor(_):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  # Embedding variable.
-  embedding_dimension = 2
 
-  with tf.Graph().as_default():
-    with hb.scope(
-        emb_backend='PAIEV', emb_device='/cpu:0', paiev_enabled={'aaa': 0}):
-      # Build columns.
-      categorical_column = tf.feature_column.categorical_column_with_identity(
-        'aaa', num_buckets=10)
-      emb_col = tf.feature_column.embedding_column(
-        categorical_column,
-        embedding_dimension,
-        initializer=tf.ones_initializer(tf.float32),
-        combiner='mean')
-
-      # Provide sparse input and get dense result.
-      sparse_input = tf.sparse.SparseTensor(
-        values=[2, 0, 1, 1],
-        indices=[[0, 0], [1, 0], [1, 4], [3, 0]],
-        dense_shape=[4, 5])
-      embedding_lookup = hb.keras.layers.DenseFeatures(
-        [emb_col])({'aaa': sparse_input})
-      with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(embedding_lookup)
 
-
-def _test_get_dense_tensor_sharded(_):
+def _test_get_embedding_variable(_):
   import tensorflow as tf
 
   import hybridbackend.tensorflow as hb
 
-  embedding_dimension = 2
-
   with tf.Graph().as_default():
-    with hb.scope(emb_backend='PAIEV', emb_device='/cpu:0'):
-      categorical_column = tf.feature_column.categorical_column_with_identity(
-        'aaa', num_buckets=sys.maxsize)
-      emb_col = tf.feature_column.embedding_column(
-        categorical_column,
-        embedding_dimension,
-        initializer=tf.ones_initializer(tf.float32),
-        combiner='mean')
-      # Provide sparse input and get dense result.
-      sparse_input = tf.sparse.SparseTensor(
-        values=[2, 0, 1, 1],
-        indices=[[0, 0], [1, 0], [1, 4], [3, 0]],
-        dense_shape=[4, 5])
-
-      embedding_lookup = hb.keras.layers.DenseFeatures(
-        [emb_col])({'aaa': sparse_input})
+    with hb.scope():
+      with hb.embedding_scope():
+        with tf.device('/cpu:0'):
+          var = tf.get_embedding_variable(
+            'var_1',
+            embedding_dim=3,
+            initializer=tf.ones_initializer(tf.float32),
+            partitioner=tf.fixed_size_partitioner(num_shards=4))
+        emb = tf.nn.embedding_lookup(
+          var, tf.cast([0, 1, 2, 5, 6, -7], tf.int64))
+      fun = tf.multiply(emb, 2.0, name='multiply')
+      loss = tf.reduce_sum(fun, name='reduce_sum')
+      opt = tf.train.FtrlOptimizer(
+        0.1,
+        l1_regularization_strength=2.0,
+        l2_regularization_strength=0.00001)
+      g_v = opt.compute_gradients(loss)
+      train_op = opt.apply_gradients(g_v)
       with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(embedding_lookup)
-
-
-def _test_get_dense_tensor_with_varscope(_):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  embedding_dimension = 2
-
-  with tf.Graph().as_default():
-    with hb.scope(emb_backend='PAIEV', emb_device='/cpu:0'):
-      partitioner = tf.min_max_variable_partitioner(
-        max_partitions=2, min_slice_size=4)
-      with tf.variable_scope('test', partitioner=partitioner):
-        categorical_column = tf.feature_column.categorical_column_with_identity(
-          'aaa', num_buckets=sys.maxsize)
-        emb_col = tf.feature_column.embedding_column(
-          categorical_column,
-          embedding_dimension,
-          initializer=tf.ones_initializer(tf.float32),
-          combiner='mean')
-        # Provide sparse input and get dense result.
-        sparse_input = tf.sparse.SparseTensor(
-          values=[2, 0, 1, 1],
-          indices=[[0, 0], [1, 0], [1, 4], [3, 0]],
-          dense_shape=[4, 5])
-
-        embedding_lookup = hb.keras.layers.DenseFeatures(
-          [emb_col])({'aaa': sparse_input})
-        with tf.train.MonitoredTrainingSession('') as sess:
-          return sess.run(embedding_lookup)
-
-
-def _test_embedding_column_with_optimizer(_, lr):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(seed=42, emb_backend='PAIEV', emb_device='/cpu:0'):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'ad0', num_buckets=sys.maxsize),
-          dimension=20,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'ad1', num_buckets=sys.maxsize),
-          dimension=30,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'ad2', num_buckets=sys.maxsize),
-          dimension=40,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'user0', num_buckets=sys.maxsize),
-          dimension=20,
-          initializer=tf.ones_initializer(tf.float32)),
-      ]
-      features = {
-        'ad0': tf.constant([0, 1, 3, 2]),
-        'ad1': tf.constant([1, 5, 3, 4]),
-        'ad2': tf.constant([5, 2, 7, 4]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      out_emb = hb.keras.layers.DenseFeatures(
-        columns, num_groups=None)(features)
-      loss = tf.reduce_mean(out_emb)
-      opt = tf.train.AdagradOptimizer(lr)
-      step = tf.train.get_or_create_global_step()
-      train_op = opt.minimize(loss, global_step=step)
-
-      final_loss = None
-      with tf.train.MonitoredTrainingSession(
-          '',
-          hooks=[
-            tf.train.StopAtStepHook(last_step=100),
-            tf.train.NanTensorHook(loss),
-            tf.train.LoggingTensorHook(
-              tensors={'loss': loss, 'step': step},
-              every_n_iter=20)]) as sess:
-        while not sess.should_stop():
-          final_loss = sess.run(loss)
-          sess.run(train_op)
-      return final_loss
-
-
-def _test_get_dense_tensor_disable_concat(_):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(
-        emb_enable_concat=False,
-        emb_backend='PAIEV',
-        emb_device='/cpu:0'):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'ad0', num_buckets=sys.maxsize),
-          dimension=20,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'user0', num_buckets=sys.maxsize),
-          dimension=20,
-          initializer=tf.ones_initializer(tf.float32)),
-      ]
-      features = {
-        'ad0': tf.sparse.SparseTensor(
-          values=[0, 1, 3, 2],
-          indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
-          dense_shape=[2, 2]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      embs = hb.keras.layers.DenseFeatures(columns)(features)
-
-      with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(embs)
-
-
-def _test_embedding_column_with_coalescing(_, lr):
-  os.environ['HYBRIDBACKEND_DEFAULT_COMM'] = 'NCCL'
-
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(seed=42, emb_backend='PAIEV', emb_device='/cpu:0'):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'ad0', num_buckets=sys.maxsize),
-          dimension=20,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'ad1', num_buckets=sys.maxsize),
-          dimension=30,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'ad2', num_buckets=sys.maxsize),
-          dimension=40,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            'user0', num_buckets=sys.maxsize),
-          dimension=20,
-          initializer=tf.ones_initializer(tf.float32)),
-      ]
-      features = {
-        'ad0': tf.constant([0, 1, 3, 2]),
-        'ad1': tf.constant([1, 5, 3, 4]),
-        'ad2': tf.constant([5, 2, 7, 4]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      out_emb = hb.keras.layers.DenseFeatures(columns, num_groups=2)(features)
-      loss = tf.reduce_mean(out_emb)
-      opt = tf.train.AdagradOptimizer(lr)
-      step = tf.train.get_or_create_global_step()
-      train_op = opt.minimize(loss, global_step=step)
-
-      final_loss = None
-      with tf.train.MonitoredTrainingSession(
-          '',
-          hooks=[
-            tf.train.StopAtStepHook(last_step=100),
-            tf.train.NanTensorHook(loss),
-            tf.train.LoggingTensorHook(
-              tensors={'loss': loss, 'step': step},
-              every_n_iter=20)]) as sess:
-        while not sess.should_stop():
-          final_loss = sess.run(loss)
-          sess.run(train_op)
-      return final_loss
-
-
-def _test_embedding_column_with_function(_, lr):
-  os.environ['HYBRIDBACKEND_DEFAULT_COMM'] = 'NCCL'
-
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  def train_fn():
-    columns = [
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'ad0', num_buckets=sys.maxsize),
-        dimension=20,
-        initializer=tf.ones_initializer(tf.float32)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'ad1', num_buckets=sys.maxsize),
-        dimension=30,
-        initializer=tf.ones_initializer(tf.float32)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'ad2', num_buckets=sys.maxsize),
-        dimension=40,
-        initializer=tf.ones_initializer(tf.float32)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'user0',
-          num_buckets=sys.maxsize),
-        dimension=20,
-        initializer=tf.ones_initializer(tf.float32)),
-    ]
-    features = {
-      'ad0': tf.constant([0, 1, 3, 2]),
-      'ad1': tf.constant([1, 5, 3, 4]),
-      'ad2': tf.constant([5, 2, 7, 4]),
-      'user0': tf.constant([2, 5, 4, 7])
-    }
-    out_emb = hb.keras.layers.DenseFeatures(columns)(features)
-    loss = tf.reduce_mean(out_emb)
-    opt = tf.train.AdagradOptimizer(lr)
-    step = tf.train.get_or_create_global_step()
-    return loss, opt.minimize(loss, global_step=step)
-
-  with hb.scope(
-      seed=42,
-      emb_num_groups=2,
-      emb_backend='PAIEV',
-      emb_device='/cpu:0'):
-    loss, train_op = train_fn()
-    final_loss = None
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=[
-          tf.train.StopAtStepHook(last_step=100),
-          tf.train.NanTensorHook(loss),
-          tf.train.LoggingTensorHook(
-            tensors={'loss': loss},
-            every_n_iter=20)]) as sess:
-      while not sess.should_stop():
-        final_loss = sess.run(loss)
-        sess.run(train_op)
-    return final_loss
-
-
-def _test_embedding_column_with_function_unique(_, lr):
-  os.environ['HYBRIDBACKEND_DEFAULT_COMM'] = 'NCCL'
-
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  def train_fn():
-    columns = [
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'ad0', num_buckets=sys.maxsize),
-        dimension=20,
-        initializer=tf.ones_initializer(tf.float32)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'ad1', num_buckets=sys.maxsize),
-        dimension=30,
-        initializer=tf.ones_initializer(tf.float32)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'ad2', num_buckets=sys.maxsize),
-        dimension=40,
-        initializer=tf.ones_initializer(tf.float32)),
-      tf.feature_column.embedding_column(
-        tf.feature_column.categorical_column_with_identity(
-          'user0',
-          num_buckets=sys.maxsize),
-        dimension=20,
-        initializer=tf.ones_initializer(tf.float32)),
-    ]
-    features = {
-      'ad0': tf.constant([0, 1, 3, 2]),
-      'ad1': tf.constant([1, 5, 3, 4]),
-      'ad2': tf.constant([5, 2, 7, 4]),
-      'user0': tf.constant([2, 5, 4, 7])
-    }
-    out_emb = hb.keras.layers.DenseFeatures(columns)(features)
-    loss = tf.reduce_mean(out_emb)
-    opt = tf.train.AdagradOptimizer(lr)
-    step = tf.train.get_or_create_global_step()
-    return loss, opt.minimize(loss, global_step=step)
-
-  with hb.scope(
-      seed=42,
-      emb_num_groups=2,
-      emb_backend='PAIEV',
-      emb_device='/cpu:0',
-      emb_unique={'ad0': True}):
-    loss, train_op = train_fn()
-    final_loss = None
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=[
-          tf.train.StopAtStepHook(last_step=100),
-          tf.train.NanTensorHook(loss),
-          tf.train.LoggingTensorHook(
-            tensors={'loss': loss},
-            every_n_iter=20)]) as sess:
-      while not sess.should_stop():
-        final_loss = sess.run(loss)
-        sess.run(train_op)
-    return final_loss
-
-
-def _test_get_dense_tensor_with_segment_rank(_):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  embedding_dimension = 2
-
-  def lookup_fn():
-    sparse_input = tf.sparse.SparseTensor(
-      values=[2, 0, 1, 1],
-      indices=[[0, 1, 1], [0, 1, 2], [1, 1, 1], [1, 1, 2]],
-      dense_shape=[4, 2, 3])
-    categorical_column = tf.feature_column.categorical_column_with_identity(
-      'aaa', num_buckets=sys.maxsize)
-    emb_col = tf.feature_column.embedding_column(
-      categorical_column,
-      embedding_dimension,
-      initializer=tf.ones_initializer(tf.float32),
-      combiner='mean')
-    return hb.keras.layers.DenseFeatures([emb_col])({'aaa': sparse_input})
-
-  with tf.Graph().as_default():
-    with hb.scope(
-        emb_segment_rank={'aaa': 1},
-        emb_backend='PAIEV',
-        emb_device='/cpu:0'):
-      embs = lookup_fn()
-      with tf.train.MonitoredTrainingSession('') as sess:
-        return sess.run(embs)
-
-
-def _test_shared_embedding_column(_, lr):
-  import tensorflow as tf
-
-  import hybridbackend.tensorflow as hb
-
-  with tf.Graph().as_default():
-    with hb.scope(seed=42, emb_backend='PAIEV', emb_device='/cpu:0'):
-      columns = [
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad1', num_buckets=sys.maxsize),
-          dimension=30,
-          initializer=tf.ones_initializer(tf.float32)),
-        tf.feature_column.embedding_column(
-          tf.feature_column.categorical_column_with_identity(
-            key='ad2', num_buckets=sys.maxsize),
-          dimension=40,
-          initializer=tf.ones_initializer(tf.float32)),
-      ]
-      columns += tf.feature_column.shared_embedding_columns(
-        [tf.feature_column.categorical_column_with_identity(
-          key='ad0', num_buckets=sys.maxsize),
-         tf.feature_column.categorical_column_with_identity(
-           key='user0', num_buckets=sys.maxsize)],
-        dimension=20,
-        initializer=tf.ones_initializer(tf.float32))
-      features = {
-        'ad0': tf.constant([0, 1, 3, 2]),
-        'ad1': tf.constant([1, 5, 3, 4]),
-        'ad2': tf.constant([5, 2, 7, 4]),
-        'user0': tf.constant([2, 5, 4, 7])
-      }
-      out_emb = hb.keras.layers.DenseFeatures(columns)(features)
-      loss = tf.reduce_mean(out_emb)
-      opt = tf.train.AdagradOptimizer(lr)
-      step = tf.train.get_or_create_global_step()
-      train_op = opt.minimize(loss, global_step=step)
-
-      final_loss = None
-      with tf.train.MonitoredTrainingSession(
-          '',
-          hooks=[
-            tf.train.StopAtStepHook(last_step=100),
-            tf.train.NanTensorHook(loss),
-            tf.train.LoggingTensorHook(
-              tensors={'loss': loss, 'step': step},
-              every_n_iter=20)]) as sess:
-        while not sess.should_stop():
-          final_loss = sess.run(loss)
-          sess.run(train_op)
-      return final_loss
+        emb_result, loss_result, _ = sess.run([emb, loss, train_op])
+        return (emb_result, loss_result)
 
 
 @unittest.skipUnless(
@@ -496,101 +70,13 @@ class DeepRecEVTest(unittest.TestCase):
   def setUp(self):  # pylint: disable=invalid-name
     os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
 
-  def test_get_dense_tensor(self):
-    results = hbtest.Spawn()(_test_get_dense_tensor)
-    np.testing.assert_allclose(
-      results[0],
-      [[1., 1.],
-       [1., 1.],
-       [0., 0.],
-       [1., 1.]], rtol=1e-6)
-
-  def test_get_dense_tensor_sharded(self):
-    results = hbtest.Spawn(2)(_test_get_dense_tensor_sharded)
-    np.testing.assert_allclose(
-      results[0],
-      [[1., 1.],
-       [1., 1.],
-       [0., 0.],
-       [1., 1.]], rtol=1e-6)
-    np.testing.assert_allclose(
-      results[1],
-      [[1., 1.],
-       [1., 1.],
-       [0., 0.],
-       [1., 1.]], rtol=1e-6)
-
-  def test_get_dense_tensor_with_varscope(self):
-    results = hbtest.Spawn(2)(_test_get_dense_tensor_with_varscope)
-    np.testing.assert_allclose(
-      results[0],
-      [[1., 1.],
-       [1., 1.],
-       [0., 0.],
-       [1., 1.]], rtol=1e-6)
-    np.testing.assert_allclose(
-      results[1],
-      [[1., 1.],
-       [1., 1.],
-       [0., 0.],
-       [1., 1.]], rtol=1e-6)
-
-  def test_embedding_column_with_optimizer(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_optimizer(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.999929, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.999929, rtol=1e-6)
-
-  def test_get_dense_tensor_disable_concat(self):
-    results = hbtest.Spawn()(_test_get_dense_tensor_disable_concat)
-    np.testing.assert_equal(len(results[0]), 2)
-
-  def test_embedding_column_with_coalescing(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_coalescing(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.999929, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.999929, rtol=1e-6)
-
-  def test_embedding_column_function(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_function(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.999929, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.999929, rtol=1e-6)
-
-  def test_embedding_column_function_unique(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_embedding_column_with_function_unique(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.999929, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.999929, rtol=1e-6)
-
-  def test_get_dense_tensor_with_segment_rank(self):
-    results = hbtest.Spawn(2)(_test_get_dense_tensor_with_segment_rank)
-    np.testing.assert_allclose(
-      results[0],
-      [[0., 0.],
-       [1., 1.],
-       [0., 0.],
-       [1., 1.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.]], rtol=1e-6)
-    np.testing.assert_allclose(
-      results[1],
-      [[0., 0.],
-       [1., 1.],
-       [0., 0.],
-       [1., 1.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.],
-       [0., 0.]], rtol=1e-6)
+  def test_get_embedding_variable(self):
+    results = hbtest.Spawn()(_test_get_embedding_variable)
+    print(results)
 
-  def test_shared_embedding_column(self):
-    results = hbtest.Spawn(2)(
-      lambda rank: _test_shared_embedding_column(rank, 0.0001))
-    np.testing.assert_allclose(results[0], 0.999923, rtol=1e-6)
-    np.testing.assert_allclose(results[1], 0.999923, rtol=1e-6)
+  def test_get_embedding_variable_2g(self):
+    results = hbtest.Spawn(2)(_test_get_embedding_variable)
+    print(results)
 
 
 # pylint: enable=missing-docstring