김태수, DietCode: Automatic optimization for dynamic tensor program

[tae-su-kim]

11/1 발표에 대한 질문 답변입니다!

[tae-su-kim]

Q: 설명해 주신 연구들이 주로 모델 deploy에서 최적화 하는 것이 맞지요? 혹시 학습시에는 최적화 하거나 하는 연구는 따로 없을까요?
A: 네 제가 설명드린 논문들은 주로 학습이 완료된 모델을 특정 하드웨어에 deploy 하는 경우를 다루고 있습니다. 제가 관심을 가지고 있는 분야가 주로 inference여서 학습 관련 지식은 부족합니다만, 지금 생각나는 페이퍼는 언젠가 학회 워크샵에서 봤던 Song han 교수님의 gradient pruning 관련 논문입니다. 아마 gradient pruning / gradient quantization 등을 통해 여러 device를 사용할때 gradient traffic을 줄여서 속도 개선을 꾀하는 논문들이 있을 것 같습니다.

[sunggg]

Training시에 automatic optimization하는 Astra라는 work도 관련 있을 것 같습니다.

[tae-su-kim]

Q: onnx 변환을 하게될 경우에 특정 연산의 경우는 변환이 되지 않아 지원이 안 되는 경우도 있는데, autoTVM, ansor, Dietcode 같은 방법들에서도 특정 연산이 들어감으로 인해서 지원이 안 되는 경우들이 발생하는지 궁금합니다. R18이 가능하다고 했는데, 여기에 사용되는 Conv가 deformable convolution 같은 것으로 대체되든지 할 경우에는 적용이 안 된다거나, 결국 어느 정도의 범위의 모델에서 사용이 가능한 것인지 궁금합니다.
A: ONNX에서 지원 안되는 operator들이 굉장히 많죠.. 저희도 많이 고통받고 있습니다. 일단 autoTVM의 경우 template이 있으면, ansor의 경우 tensor expression으로 표현한 코드가 있으면 optimization이 가능합니다. DietCode는 일단 최적화하는 부분은 ansor와 비슷하구요. 그래서 ansor 기준으로는 다양한 operator들에 편하게 사용할 수 있을 것이라고 생각합니다. 검색해보니 deformable convolution도 TVM relay에 구현이 되어있어서 사용이 가능 해 보입니다. 만약 on-the-fly로 offset이 바뀐다면 문제가 많겠지만 학습 과정에서 offset이 고정되면 결국 deploy 시에는 특정 convolution으로 고정되니 optimization에 문제가 없을 것 같습니다.
개인적으로는 사용 가능한 연산의 범위는 on-the-fly로 memory access pattern이 바뀌거나 하는 연산이 아니라면 충분히 optimization이 가능하지 않을까라고 생각하고 있습니다. TVM IR에 구현되어 있는 operator list는 (https://tvm.apache.org/docs/reference/api/python/relay/nn.html) 를 참조하시면 되겠습니다.

[tae-su-kim]

Q: Dietcode외에 최근에 관련 논문 두 개가 나왔다고 하셨는데, 혹시 이번에 소개하지 않은 paper의 접근 방식과 Dietcode의 접근 방식의 가장 큰 차이점을 간략하게 설명해주실 수 있을까요?
A: 제가 언급했던 다른 논문 두개는 TensorIR (https://arxiv.org/pdf/2207.04296.pdf)과 Metaschedule (https://arxiv.org/pdf/2205.13603.pdf) 입니다. TensorIR은 TVM의 기존 IR 구조(TIR - schedule 이원화)를 수정해서 optimization 과정에서 TIR을 직접 수정할 수 있도록 하여 scheduling 과정을 좀 더 단순하고 강력하게 만들고, tensorization을 더 효율적으로 지원하기 위해 제안된 새로운 IR입니다. RFC가 올라왔던 것은 2020년이었던 것으로 기억하는데, 논문은 tensorization 지원 개선과 이를 바탕으로 한 auto-optimzation으로 얻은 성능 개선을 위주로 작성된 것으로 보입니다. Metaschedule은 기존 autoTVM과 Ansor를 바탕으로 HW expert가 제안하는 최적화된 특정 kernel과 autoTVM용 template, 그리고 Ansor 기반으로 auto-optimize 된 kernel을 모두 합친 큰 search space를 만들고, 이 search space에서 optimization을 수행해서 더 좋은 kernel을 찾고자 하는 연구입니다.
TensorIR과 Metaschedule은 최적화 과정에서 더 좋고 큰 search space를 만드는데 중점을 두었고, DietCode는 기존에 지원되지 않던 dynamic tensor shape를 지원하는데 중점을 둔 페이퍼이므로 결이 조금 다르다고 볼 수 있을 것 같습니다. 관심이 있으시면 TensorIR과 Metaschedule 페이퍼는 순서대로 읽어보시면 좋을 것 같습니다! 두 내용을 같이 읽으면 좋은데.. 그래서 분량상 소개를 못했습니다ㅠ

[tae-su-kim]

Q: 논문에는 Dietcode 적용 시, 시간절약에 대한 결과는 많은데, 성능이 어떻게 나왔는지에 대해서는 설명이 보이지 않는 것 같은데요. 이건 성능이 많이 떨어지지는 않을까요?
A: 넵 기본적으로 deploy 과정에서 model이 fix된 후를 가정하고 있다 보니 모델 성능에서는 차이가 발생하지 않습니다. DietCode의 경우 zero-padding을 통해 연산 과정에서 micro-kernel에 fit하지 않는 부분을 모두 zero-out 시키므로 input / output integrity에 문제가 생기지 않는다고 보시면 됩니다.

[tae-su-kim]

Q: 저는 TF와 JAX를 쓰다보니 XLA를 많이 사용하는 편인지라 관련하여 몇 가지 질문이 있습니다.

1. jax는 xla를 통해 다양한 장치에서 빠르게 동작하는 것을 강점으로 내세우고 있습니다. tvm도 비슷한 역할을 하는 것으로 보이는 데, 더 범용성이 넓어보입니다. 그렇다면 tvm도 model compiler로서 다양한 장치에서 학습이 가능하도록 만드는 중간다리가 될 수 있을까요?
2. xla의 compile속도는 무시할정도는 아니지만 대부분 수 분내로 끝나 부담이 되지 않는데, 제가 잘못들었을지 모르지만 tvm에서는 시간단위로 이야기하시는걸 보고 놀랐습니다.어떤 이유로 더 시간이 오래걸리게 되는 지 궁금합니다."

A: 제가 pytorch를 주로 사용해서 XLA에는 지식이 부족한 편인데요.. 최대한 답변을 드려보도록 하겠습니다.

1. 아마 TVM은 아직까지 inference only로 알고 있는데요, backward path가 전부 구현되어 있지는 않아서 그런 것으로 알고 있습니다. 다만 몇몇 project에서 forward path 구조가 fix되어 있는 경우 forward는 TVM 기반으로 하고, backward는 pytorch 등 다른 framework로 수행한 것을 본 적이 있는데, 이런 방식으로는 접근이 가능할 것 같습니다.
2. 기본적으로 XLA는 graph compiler이고, XLA를 사용하면 HLO에서 high-level optimization, 주로 op fusing을 수행하고 optimized graph를 CPU / GPU의 경우 LLVM IR로 변환해서 LLVM을 이용해 JIT든 AoT든 compile을 수행하는 것으로 알고 있습니다.
   TVM과 차이점은 IR - Compile 과정에서 볼 수 있습니다. 예를 들어 CONV-BN-ReLU operator가 주어진다면 XLA의 경우 IR에서 compile이 일어날 때 LLVM backend에 구현되어있는 특정 CONV-BN-ReLU template을 이용해 바로 compile이 이루어지지만, TVM의 경우 CONV-BN-ReLU에서 CONV의 경우 6 nested loop를 어떻게 tiling, vectorization 등을 통해 어떤 parameter로 최적화 할 것인지 많은 수를 테스트해보고 최종적인 machine code를 생성하기 때문에 테스트 과정에서의 시간 소요가 큰 차이를 가져오는 것으로 이해하시면 되겠습니다.

[sseung0703]

상세한 답변 감사합니다!

[tae-su-kim]

Q: Latency가 search space에 있는 인자들에 의해 미분이 가능한가요? 강화학습이나 GA의 경우 미분이 불가능한 reward 혹은 score를 통해서도 학습이 될 수 있는데, XGBOOST와 같은 ML 알고리즘들도 미분이 불가능한 objective를 학습할 수 있는지 혹은 latency가 미분이 가능한지 궁금합니다.

A: Latency는 측정값이므로 미분 불가능합니다. latency를 kernel의 tunable option에 대해 미분할 수 있다면 gradient based search가 가능해지겠지만.. 이는 불가능하므로 AutoTVM의 XGBoost Tuner에서는 latency를 예측하는 loss (regression loss) 혹은 kernel의 latency 크기를 pairwise로 비교하는 loss (rank loss)를 정의해 학습하고 있습니다.

[tae-su-kim]

Q: 4bit quantized kernel 코드 관련해서 tvm에 작성된 경로를 알 수 있을까요? Nvidia tensor core에 맞춰서 작성된 것일까요? 4bit을 지원하는 범용적인 하드웨어가 tvm에서 support되는 것이 있는지도 궁금합니다.

A: Relay에 bitserial_conv2d와 bitserial_dense operator가 정의되어 있습니다 (https://tvm.apache.org/docs/reference/api/python/relay/nn.html#tvm.relay.nn.bitserial_conv2d). 4-bit Tensor Core 기준은 아니고, CPU 등에서 bitserial computation을 통해 sub-4-bit conv2d / dense operation을 구현하도록 되어 있습니다. 제가 아는 한에서는 TVM stack에 아직 4 bit Tensor Core 지원이 구현되지 않은 상태이므로 commodity hardware 기준으로는 아직 지원되지 않는 것으로 봐야 할 것 같습니다.

[tae-su-kim]

Q: 모델 latency 등을 직접 측정해보면 accuracy와는 다르게 다양한 이유로 측정할 때마다 variation이 생기던데, autoTVM 등으로 tuning할때는 작은 기본 연산들로만 측정해서 variation 영향이 작은 편일까요? 만약 편차가 있다면 측정 latency결과를 cost model에 넣어서 예측할 때 어떻게 잘 반영할 수 있을지 궁금합니다.

A: 실제로 휴대폰 등에서 benchmarking을 해보면 발열과 background job 등으로 인해 variation이 큰 편인데요, 일반적으로 이를 해결하기 위해 같은 kernel을 여러번 run하고 평균을 내는 식으로 접근합니다. AutoTVM의 기본 옵션은 1회 inference로 warm-up 수행 후 4회 inference의 평균치를 다시 3번 평균 내는 식으로 되어 있습니다 (https://tvm.apache.org/docs/reference/api/python/autotvm.html#tvm.autotvm.measure.measure_methods.RPCRunner).

[tae-su-kim]

Q: 이러한 Compiler 단의 Design Space 를 최적화하고 Evaluation 을 할 때, 실제 Edge Device 에서 Deploy 하는 상황에서 성능 비교하는 경우도 있는지 궁금하고, Edge Device Inference 에서의 Compiler 최적화 효과와 GPGPU 에서의 Compiler 최적화 효과의 정도가 얼마나 다를지 발표자님의 견해를 들어보고 싶습니다.

A: AutoTVM의 경우 microTVM이라는 micro-controller 전용 extension이 있어 edge device에서의 deploy도 가능합니다. 아직 hardware 지원이 다양하지 않은 편인 한계가 있습니다만, 기본 틀은 제공이 되므로 필요한 경우 유저가 확장하여 성능 비교와 optimization을 수행할 수 있을 것이라고 생각합니다.
Edge Device와 GPGPU간의 비교에 대한 견해는 간단하게 정리해서 말씀드리기엔 내용이 너무 많은 것 같습니다..ㅎㅎ 큰 틀에서는 edge의 경우 (특히 microcontroller의 경우) device 별 instruction set 간의 차이나 memory size, memory bandwidth 등에 영향을 크게 받는 것 같습니다. 저는 이런 측면에서 kernel auto-optimization을 통한 memory access pattern의 개선 등이 큰 이득을 가져올 수 있다고 보는 편입니다. 이쪽은 Song Han 교수님의 MCUNet에서 출발하는 여러 논문들이 restricted memory 상에서 효율적인 inference를 위한 연구를 하고 있으니 이 부분을 참고해보시면 좋을 것 같습니다.
반면에 GPGPU는 NVIDIA GPU와 CUDA 기반 acceleration으로 한정하여 생각해 보면 기본적으로 NVIDIA에서 워낙 vendor library를 잘 제공하고 있다 보니 TVM 등의 외부 framework를 이용한 최적화는 좀 제한적인 부분이 있다고 생각합니다. 특히 TensorRT가 graph-level optimization도 포함하고 있고, 알려진 operator들에 대해서는 CuBLAS, CuDNN, CUTLASS 가리지 않고 다 비교해서 좋은 커널을 뽑아 쓰도록 해주기 때문에.. operator 형태가 정말 skew되어 있는 경우가 아니라면 기본적으로는 TensorRT 쓰는게 베스트일 것이라고 생각합니다. TensorRT도 일종의 compiler 최적화라고 본다면 그냥 CuDNN 사용하는 것 보다는 성능 개선이 크기 때문에 어쨌든 GPU computing에서도 compiler 최적화 효과가 충분히 보인다라고 할 수도 있겠네요.

[tae-su-kim]

Q: "Ansor 설명을 해주실 때, batch size가 커질수록 다른 프레임워크와 비교했을때 성능차이가 좁혀지는 데 어째서 이런현상이 발생하는 걸까요?
배치가 커질수록 최적화가 어렵다..라는 뜻으로 해석해야 하는걸까요? 아니면 다른 이유가 있을까요?"

A: Batch size가 커질수록 workload가 커지게 되고, vendor library 설계시 가정하는 일반적인 tensor size에 더 잘 맞기 때문에 vendor library의 성능이 올라간다라고 해석하는 것이 맞아 보입니다. batch가 작아질수록 일종의 skewed operator가 되는 것이고 (GPU에서 모델이 엄청 크거나 한 이유가 아니라면 batch=1 쓸 일이 많지 않기 때문에) vendor library는 이러한 경우를 잘 지원하지 못하지만 ansor는 skewed operator에 대해서도 최적화를 수행할 수 있기 때문에 성능 개선 폭이 크다고 보면 되겠습니다.

[tae-su-kim]

Q: SqueezeBits에서도 dynamic model에 관심이 많이 있으신가요?

A: 최근에 이슈가 되고 있는 모델 중 많은 부분이 dynamic shape tensor를 포함하고 있기 때문에 저희도 해당 연산들을 효율적으로 support할 수 있도록 준비하고 있습니다. 아마 앞으로 많은 시간을 이쪽 최적화에 사용할 것 같습니다!

[tae-su-kim]

Q: 아직까지 기존 라이브러리를 사용한 것에 비해 큰 성능 개선이 없는 것으로 보입니다. 단순 성능보다 오히려 AITemplate처럼 다양한 하드웨어에 배포하거나 nvFuser 처럼 high level program을 자동으로 fusion하는 방법이 더 유용할 수 있다고 생각되는데 이 부분에 대해 생각이 궁금합니다.

A: 사실 저도 아직까지는 DL compiler를 쓰는것이 vendor library를 사용하는 경우에 비해 특정 케이스가 아니면 이득이 많지 않다고 생각하는 편입니다. 말씀해 주신 대로 다양한 hardware로의 배포를 지원하도록 하는 IR의 구현과 graph-level optimization을 잘 구현해 놓는 것이 지금 당장의 deployment에 있어서는 확실히 중요한 문제라는 점에 공감하구요. 사실 TVM framework가 자체 IR과 저런 optimization 기능을 구현하고 있기 때문에 이 부분은 꽤 유용하게 사용할 수 있어서 현재로서는 auto optimzation 기능보다 이쪽이 TVM framework 자체의 메리트라고 생각합니다. 다만 미래를 생각할 때 DL compiler가 고도화되는 것이 (특히 제 밥줄에) 대단히 중요한 문제이기 때문에 어떻게 하면 성능을 개선할 수 있을지 고민하는 방향으로 논문들을 읽고 있는 편입니다.

[tae-su-kim]

Q: SageMaker, OctoML과 같은 DL 컴파일러 기술이 적용된 서비스가 최근에 많이 사용되고 있습니다. 재직하시는 회사에서도 개발한 기술이 이러한 서비스로 제공이 예정되어 있는지 알고 싶습니다.

A: 저희도 model compression에 쉽게 접근하고 실제로 mobile 환경이나 GPU 등에서 경량화된 model을 쉽게 test / deploy 할 수 있도록 하는 서비스를 현재 개발중입니다. 빠른 시일 내로 만나보실 수 있을 예정입니다!