NCCL을 이용한 Efficient한 Tensorflow MultiGPU Training 코드 작성하기

tensorflow

(Curtis Kim) #1

예전에 Tensorpack과 Multigpu를 활용한 빠른 트레이닝 코드 작성하기에서 언급한 것처럼, multi gpu를 이용해서 트레이닝 코드를 작성하면 당연히 효율이 좋아집니다. 이 때, input queue가 충분히 빠르게 채워져야하는 것이나 네트워크 그래프가 효과적으로 잘 작성되어야하는 것 등이 필수적입니다. 그렇지 않으면, 오히려 1개의 gpu를 사용했을 때보다 성능이 잘 안나오는 경우가 있기도 합니다.

이 글에서는 multi gpu를 더욱 효과적으로 사용하기 위해 NCCL:NVIDIA Collective Communications Library 를 사용하는 법을 간단히 소개합니다.

핵심 3줄 요약

  • Weight 를 모든 GPU에 싱크하는 형태는 비효율적. 연산할 때마다 Weight를 다른 곳에서 복사한 후 읽어서 처리하는 구조는 느릴 수 밖에 없다. 따라서 Weight는 모든 GPU에서 각자 업데이트한다.
  • 동일한 Gradient를 모든 GPU에 전송하는 것이 Weight를 보내는 것보다 빠르다. 초기 값이 동일한 상태에서 반복적으로 동일한 Gradient를 적용하면 동일한 Weight가 유지된다.
  • NCCL의 all-reduce sum 을 이용하면, gradient sum을 여러 gpu 간에 효과적으로 할 수 있고, 이를 이용해서 multi gpu 트레이닝 코드를 작성하는 것이 상대적으로 더 효과적일 수 있다.

Backgrounds

텐서플로우 관련해서 위와 같은 글처럼, multi-gpu 에서의 퍼포먼스에 대한 문의가 꽤 됩니다. 위 글에서는 gpu 수를 늘려보면서 실험해봤더니 처리량이 늘긴하지만, 기대했던 것만큼 늘지는 않았다는 게 핵심입니다.

image

GPU 8장을 꼳으면 기대하는 것이 GPU 1장 대비 8배 향상을 원하는데 위 테이블을 보면, 대략 절반에도 미치지 못하네요.

Tensorflow official Benchmark

image

https://www.tensorflow.org/performance/benchmarks

하지만 텐서플로우 공식 벤치마크에서는, 거의 이상에 가까운 그래프를 보여주고 있습니다. 그래서 텐서플로우에서 multi gpu를 사용하는 공식 예제를 먼저 살펴보는게 좋겠네요.

Tensorflow CIFAR 10 Multi-GPU Code

예제코드를 살펴보면 핵심이 되는 부분은 아래입니다.


    # Calculate the gradients for each model tower.
    tower_grads = []
    with tf.variable_scope(tf.get_variable_scope()):
      for i in xrange(FLAGS.num_gpus):
        with tf.device('/gpu:%d' % i):
          with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:
            # Dequeues one batch for the GPU
            image_batch, label_batch = batch_queue.dequeue()
            # Calculate the loss for one tower of the CIFAR model. This function
            # constructs the entire CIFAR model but shares the variables across
            # all towers.
            loss = tower_loss(scope, image_batch, label_batch)

            # Reuse variables for the next tower.
            tf.get_variable_scope().reuse_variables()

            # Retain the summaries from the final tower.
            summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)

            # Calculate the gradients for the batch of data on this CIFAR tower.
            grads = opt.compute_gradients(loss)

            # Keep track of the gradients across all towers.
            tower_grads.append(grads)

    # We must calculate the mean of each gradient. Note that this is the
    # synchronization point across all towers.
    grads = average_gradients(tower_grads)

위 코드가 하는 일은 동일한 네트워크 그래프를 GPU 수 만큼 만들고, 동일한 Weight Variable을 공유하도록 한 것입니다. 즉 네트워크는 GPU 수 만큼 독립적으로 있습니다만, Weight 값은 1개인 상황입니다. 각 GPU에서 서로 다른 데이터를 이용해 Gradient를 구하고 이를 average_gradients라는 함수를 이용해서 합치고 있습니다.

def average_gradients(tower_grads):
  """Calculate the average gradient for each shared variable across all towers.
  Note that this function provides a synchronization point across all towers.
  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over individual gradients. The inner list is over the gradient
      calculation for each tower.
  Returns:
     List of pairs of (gradient, variable) where the gradient has been averaged
     across all towers.
  """
  average_grads = []
  for grad_and_vars in zip(*tower_grads):
    # Note that each grad_and_vars looks like the following:
    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
    grads = []
    for g, _ in grad_and_vars:
      # Add 0 dimension to the gradients to represent the tower.
      expanded_g = tf.expand_dims(g, 0)

      # Append on a 'tower' dimension which we will average over below.
      grads.append(expanded_g)

    # Average over the 'tower' dimension.
    grad = tf.concat(axis=0, values=grads)
    grad = tf.reduce_mean(grad, 0)

    # Keep in mind that the Variables are redundant because they are shared
    # across towers. So .. we will just return the first tower's pointer to
    # the Variable.
    v = grad_and_vars[0][1]
    grad_and_var = (grad, v)
    average_grads.append(grad_and_var)
  return average_grads

average_gradients는 각 gpu에서 얻어진 gradient를 평균 내는 역할을 합니다. 즉 동일한 변수에 대해 서로 다르게 구해진 gradient 값을 서로 평균내는 것입니다.

이렇게 얻어진 gradient를 Weight에 적용해 Backpropagation을 하게 됩니다.

문제는, 이 코드조차도, 일반적인 상황에서 돌려보면 gpu 수가 늘어나는 것만큼 성능이 늘어나지 않는다는 데 있습니다.

문제는 무엇일까

Multi GPU를 사용하는 방법에는 여러가지 방법이 있습니다만, Parameter Server를 사용하지 않고, 하나의 로컬 머신에서 처리하는 위와 같은 예제에서 속도 저하가 생기는 요인은

다른 GPU를 기다리는 시간만큼 자원을 사용하지 않게 된다.

는 점입니다. 이 부분을 중심으로 위 코드를 프로파일링해보면, 이런 문제가 있습니다.

  1. 다른 GPU에서 Gradient 계산이 끝날 때까지 기다려야 한다 (아주 큰 영향일 것 같지는 않습니다)
  2. 모든 GPU에서 Gradient 계산이 끝나면, 이를 CPU 또는 특정 GPU로 모아서, 평균을 구한다. 그리고 Weight 업데이트를 한다.
  3. 트레이닝 할 때 Weight 값은 특정 GPU(아마 첫번째 GPU)에서 매번 읽어온다.

(1) 의 경우에는 Async한 Weight Update를 하도록 변경하면 가능합니다만, 여기에서는 다루지 않겠습니다.

(2)와 (3)에 대해서만 해결책을 생각해보도록 하겠습니다. (2)와 (3)을 요약해서 이야기하면, 통신 비용이 발생하는 형태의
Gradient Update와 Weight Reading 연산이라는 문제입니다.

NCCL 을 이용해서 새롭게 작성하기

이 코드는 주로 Tensorpack 내에서 이곳저곳을 참조해서 작성했습니다.

NCCL은 NVidia에서 여러 GPU 간에 주로 행해지는 몇가지 연산을 지원하기 위해 만들어진 패키지입니다. 예를 들어서 여러 GPU에 있는 값을 읽어서 덧셈을 하거나, 값을 전달하는 Broadcasting 등을 수행하는 등에서 좋은 성능을 내기 위해 특별히 제공되고 있습니다.

Tensorflow에서도 NCCL에 대한 기본적인 지원을 하고 있는데, 아래의 문서를 참고하시면 됩니다.

https://www.tensorflow.org/api_docs/python/tf/contrib/nccl

저는 이 중에서 all_sum 이라는 함수를 이용해 구현했습니다. all_sum 은 all-reduce sum 라는 뜻입니다. 즉, 각 gpu에서 값을 읽어 합한 후, 그 값을 모든 gpu에 다시 저장하는 함수입니다. 이 함수를 이용해서 작성한 예제를 살펴보겠습니다.

변해야 할 것들이 좀 있습니다.

  1. 트레이닝할 때 Weight 값을 특정 GPU에서 읽어오던 기존 방식에서 벗어나기 위해서, 모든 GPU에 동일한 Weight를 유지한다.
  2. NCCL의 all_sum 을 이용해서 여러 GPU로부터의 Gradient의 합을 구한다.
  3. all_sum은 구해진 합을 각 GPU에 broadcast 하므로, 각 gpu의 Weight에 적용해 gradient만큼 변화를 만든다.

즉, 기존의 방식에 비해서 다른 gpu에서 데이터를 읽고, 카피하는 비용이 현저하게 줄어듭니다. 왜냐하면

  • Weight를 복사하던 기존 방식에서 Gradient 만을 전달하는 방식으로 변경됨
  • Gradient Sum을 구해 각 gpu로 전달하는 과정이 NCCL을 이용해 효율적으로 처리됨

이기 때문입니다.

            grad_list = []
            for gpu_idx in range(num_gpu):
                logger.info('creating gpu tower @ %d' % (gpu_idx + 1))
                with tf.device(tf.DeviceSpec(device_type="GPU", device_index=gpu_idx)), tf.variable_scope('tower%d' % gpu_idx):
                    logit, _ = self.__create_network_for_imagenet(
                        ph_train_image_batch[gpu_idx],
                        is_training=self.is_training,
                        is_reuse=False,                                       # **** NO SHARING WEIGHTS!! *****
                        depth_multiplier=depth_multiplier
                    )

                    grad_list.append([x for x in self.optimizer.compute_gradients(loss) if x[0] is not None])

위 부분은 기존처럼 각 gpu에 네트워크 그래프를 만드는 것인데, Weight Sharing을 하지 않습니다. 그리고 gradient를 각각 구해둔 뒤, 아래와 같은 함수를 이용해 gradient를 평균을 구합니다. 이때 사용하는 것이 nccl의 all_sum 입니다.

def allreduce_grads(all_grads, average=True):
    from tensorflow.contrib import nccl
    nr_tower = len(all_grads)
    if nr_tower == 1:
        return all_grads
    new_all_grads = []  # N x K
    for grads in zip(*all_grads):
        summed = nccl.all_sum(grads)

        grads_for_devices = []  # K
        for g in summed:
            with tf.device(g.device):
                # tensorflow/benchmarks didn't average gradients
                if average:
                    g = tf.multiply(g, 1.0 / nr_tower, name='allreduce_avg')
            grads_for_devices.append(g)
        new_all_grads.append(grads_for_devices)

    # transpose to K x N
    ret = list(zip(*new_all_grads))
    return ret

위 쪽의 CIFAR 10 예제와 유사하나

  • nccl의 all_sum을 이용한다는 점
  • all_sum으로 구해진 데이터가 각 gpu에 존재한다는 점

이 다릅니다. 참고로 각 gpu에 존재하는 데이터는 모두 동일한 gradient 입니다.

# optimizer using NCCL
train_ops = []
for idx, grad_and_vars in enumerate(grads):
    with tf.name_scope('apply_gradients'), tf.device(tf.DeviceSpec(device_type="GPU", device_index=idx)):
        # apply_gradients may create variables. Make them LOCAL_VARIABLES
        with override_to_local_variable(enable=idx > 0):
            train_ops.append(self.optimizer.apply_gradients(grad_and_vars, name='apply_grad_{}'.format(idx)))

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
    self.optimize_op = tf.group(*train_ops, name='train_op')

마지막으로는 이렇게 구해진 Gradient들을 각 GPU에서 적용하도록 처리하는 것입니다. 이렇게 함으로써 기존에 Weight를 복사하던 부분이 아예 사라지게 됩니다.

Performance?

위 코드는 제가 mobilenet-v2 를 재현하면서 작성한 multi gpu 코드입니다. mobilenet v2 설명은 (MobilenetV2) Inverted Residuals and Linear Bottlenecks: Mobile Networks for Classification, Detection and Segmentation 에서 보시면 됩니다.

image

위 그래프가 기존 GPU 8개를 사용하던 때의 사용률입니다. 절반에 미치지 못하는 사용율입니다. 또 사용율이 심하게 들쑥날쑥인 모습입니다.

image

개선된 버전입니다. 중간 중간 CPU 작업이 좀 있는 편이라 멈춤이 좀 있습니다만, 전반적으로 8 gpu를 거의 다 잘 활용한다고 볼 수 있습니다.

Tensorflow Benchmark에서처럼 이상적인 수준이 나오기 위해서는 하드웨어적으로 NVLink 같은 추가적인 구성이 있어서 gpu간 복사 비용이 매우 적거나, 아니면 이 글에서 소개하는 것처럼 코드레벨에서 오버헤드를 줄이기 위한 노력을 해야합니다.

아마 이후에는 이보다 더 개선할 수 있는 방법을 찾아 적용하고, 이를 단순한 wrapper 형태로 만들어 써보면 어떨까하는 생각을 가지고 있습니다.


(Curtis Kim) #2

인풋을 gpu 별로 서로 다른 queue를 갖게하는 것이 더 성능이 좋은 것 같습니다. tf.split을 사용해서 하나의 배치를 여러 gpu로 잘라 넣는 방식은, cpu time이 많이 걸립니다. 위 그래프를 보면 기존 대비 조금 더 개선된 것을 볼 수 있습니다.


(권동혁) #3

GPU Utilization Logging은 혹시 어떤 툴을 사용하셨는지 여쭤봐도 될까요?


(Curtis Kim) #4

저건 직접 만들어서 붙인 건데요, 주기적으로 nvidia-smi 로 가져와서 로그 쌓고 grafana 같은 시각화 도구에 붙인 겁니다. 직접 구현한 건 아니라서 … ㅎ


(권동혁) #5

아하, 네 감사합니다 :smiley:
그래프가 너무 아름다워섴ㅋ 혹시 자동화된 Tool이 있는가 해서 여쭤봤어요. ^^


(권동혁) #6

혹시 Multi-node Training은 구현해 보신 경험은 없으신가요? HOROVOD 사용을 해보려는데 ㅠㅠ 코드 찾기가 어려워서 ㅠㅠ


(Curtis Kim) #7

저도 horovod 자주 사용합니다. 예제 코드 수준만 보셔도 시작하시기엔 무리가 없을 거구, 아마 삽질은 좀 하셔야 될 것 같아요. 한 컴퓨터 내에서 멀티 gpu 하는게 아직은 현실적으로 좋은 안으로 생각합니다.


(권동혁) #8

네, 사실은 필요에 의해서 하는 거라 ㅠㅠ Slim코드를 Estimator 사용하는 코드로 고쳐가는 중입니다(Deeplab v3 코드). NVIDIA-EXAMPLE CODE를 많이 참고하는 중이기는 한데, GPU마다 다른 데이터를 잘 긁어오는지, 실제 Batch가 더 커지는 효과가 있는지, 같은 기본적인 부분부터 확인하고 가려니 막히는 게 많네요. 혹시 비슷한 작업을 진행해보신 적은 없으시죠?


(Curtis Kim) #9

제 오픈소스들이 대부분 multi gpu인데 estimator는 한번만 써봣습니다. Deepface라는 리포지터리에 올라갈 예정입니다.


(권동혁) #10

아, 네 올라오면 참고하겠습니다. 감사합니다. ^^


#11

안녕하세요, 좋은 접근 방법 정리해주셔서 감사합니다.
이 게시물의 방식대로 코드를 짜보고 있는데 한 가지 의문이 생겨서요.
예를 들어, device가 /gpu:0, /gpu:1, /gpu:2, /gpu:3 이라 했을 때 학습 전 같은 값으로 초기화는 어떤 방법으로 시켜야 하나요? 게시물만 봐서는 알기가 힘드네요.
제가 생각한 방식은 global_variable_initializer로 모든 variable을 초기화 한 후, reference device(예를 들면 /gpu:0)의 값으로 다른 device의 variable을 tf.assign을 이용해 같게 해주는 것인데요.
이 방식이 맞나요? 아니면 다른 더 효율적인 방법이 있나요?


(Curtis Kim) #12

모든 Variable에 대해서 initialize 한 후에 sync operation을 통해서 동기화 시켜줍니다. 주기적으로 동기화시켜주는 operation인데, 시작 전에 한번 호출하는 방식입니다.


#13

감사합니다! 원하던 해답이었네요!!!


(최형석) #14

안녕하세요, 글을 읽던 중 궁금한게 생겨서 여쭤봅니다.
두 가지 개선방안을 제시해주셨는데 각각의 방법이 속도를 개선하는데 어느정도의 비율을 차지하는지가 궁금합니다. 더불어 첫 번째 제안하신 variable reuse를 안하는 경우에 불이익이 없는지 궁금합니다. 감사합니다!


(Curtis Kim) #15

첫번째 방식(TF 예제)은 거의 개선이 없었고, 두번째 방식으로 해야 gpu 수에 거의 linear한 개선이 일어납니다. variable reuse는 해당 코드 아래 부분의 summary 등에서 사용되는 것으로, 실제 계산에는 영향을 미치지는 않습니다.


#16

안녕하세요, multi-gpu를 사용하려고 따라해 보았습니다. ㅎㅎ 감사합니다. 저는 랩탑에 gtx 1070 gpu 2개를 사용하는데요, nccl을 사용할때 평균 20% 의 성능(빠르기)향상이 있는 걸로 보입니다 (현재까지는요.)
다름이 아니라 저는 operation 갯수를 비교해 보았는데요, gpu 갯수가 늚에 따라 (두개의 gpu) operation과 variables 갯수도 함께 늘어서 (거의 2배) 시간이 늘어나는 부분에 대해서는 어쩔수 없다고 생각하는데 맞는 생각인가요?
그리고 위에 네트워크는 독립적으로 있다고 했지만 weight는 1개라고 하셨는데, 사실 변수는 두개로 늘어나는거죠? tower0/var0, tower1/var0 처럼요?
그리고 nccl을 사용하지 않으면 gpu0이 작동할때 gpu1이 기다린다고 하셨는데요, nccl을 사용하면 두개가 같이 돌아가나요? gpu0 돌아갈때 gpu1이 기다리면 gpu 하나만 돌리는것과 뭐가 다른가요? 오히려 말씀하셨다시피 변수들 카피하고 할당하고 이러는데 들어가는 operation 수만 더 늘어나서 불이익 아닌가 싶습니다만, 그렇게 되면 nccl을 사용하지 않는 경우 multi-gpu의 이득이 없어야 되는데 nccl을 안써도 200~300% 이득은 있는것 같네요. 어떻게 생각하시는지요? ^^

감사합니다!


(Curtis Kim) #17

네 variable은 각 gpu 메모리에 따로 보관되어야 하기 때문에 늘어나는게 맞습니다. 값이 share된다는 뜻입니다.

문제점1로 지적한 다른 gpu에서의 계산을 기다리는 것은 어차피 비슷한 연산을 각 gpu에서 하기 때문에 시간 차가 거의 없어서 큰 영향으로 보지는 않습니다. 오히려 문제점 2/3 등에서 지적한 변수 카피에서의 불이익 때문에 오래 기다린다는 뜻으로 보시면 됩니다.

감사합니다.


#18

아하! 네 제가 다른 뜻으로 이해했네요. ㅎㅎ 하나씩 따로 순서대로 돈다는 줄로 이해했습니다! 감사합니다.
질문이 하나 더 있는데요, 혹시 rnn 계열도 멀티gpu로 해보셨나요? 성능이 잘 나오는지 좀 궁금해서요 저는 rnn 계열에 시도해 봤는데 제가 코드를 잘못짯는지 15~20% 정도만 계산이 빨라지네요 ㅎㅎ 얼핏 다른 친구 얘기로는 rnn 계열에서는 잘 안된다고 들었다고 하는데 이 친구는 pytorch 사용합니다만…ㅎㅎ

감사합니다!


(Curtis Kim) #19

RNN 모델에 따라 다를 수 있을 것 같습니다. RNN은 시간 축에 따른 dependency 때문에 고려해야할 사항이 늘어날 수도 있거든요. ㅎ


#20

아 맞습니다. 모델에 따라 다를 듯 합니다.
죄송합니다만 질문 하나 더 드려도 될까요? 혹시 gpu 하나만 있는 환경에서 위의 multi gpu code를 돌려보셨나요? 왜 에러가 안나는지 모르겠어서요. gpu 갯수를 2로 해서 돌렸는데 돌아가네요… -_-ㅎㅎ 갯수 3 이상은 안돌아 가네요. 에러가 따로 나지는 않고 그냥 멈춰있네요. 없는 gpu를 찾으려고 해서 그런걸까요…
주변에 텐서 플로우를 쓰는 사람이 별로 없네요. 다들 pytorch 파네요. ㅎㅎㅎ
바쁘신데 항상 감사합니당 ㅎㅎ