TensorFlow Custom Op으로 데이터 변환 최적화하기

대규모 언어 모델용 TFRecord를 만들 때 N-gram masking 위치 계산이 전체 전처리의 절반 이상을 차지한 문제를 TensorFlow C++ Custom Op으로 해결한 사례다. Op metadata, kernel 구현, shared object 빌드, Python 로딩 과정을 코드 수준에서 설명한다. 병목은 문장당 4.697ms에서 17.147μs로 줄어 약 250배 빨라졌다.

핵심 포인트

Custom Op은 기존 TensorFlow op 조합으로 표현하기 어렵거나 비효율적인 연산, 직접 fusion할 연산에 적합하다.
BERT 입력과 NSP·SOP label을 만드는 전처리에서 확률적 N-gram masking 후보 계산이 가장 큰 병목이었다.
`REGISTER_OP`로 input·output·attribute·shape function을 선언하고 `OpKernel.Compute()`에 C++ 계산을 구현한다.
`REGISTER_KERNEL_BUILDER`로 metadata와 CPU kernel을 연결한 뒤 TensorFlow compile·link flag를 사용해 `.so`로 빌드한다.
`tf.load_op_library()`로 shared object를 읽으면 PascalCase op가 snake_case Python 함수로 노출된다.
유지보수와 디버깅 비용이 크므로 일반 op으로 충분한 경우 피하고, 측정된 핵심 병목에 제한적으로 적용해야 한다.

상세 정리

상황: 방대한 corpus를 TPU 학습용 TFRecord로 변환해야 했다. 결과에는 BERT의 input ids·attention mask, masking 위치, NSP 또는 SOP label이 들어간다.
동적 masking: 여러 epoch마다 다른 위치를 가리기 위해 실제 masked input을 저장하지 않고 masking position만 기록했다. ALBERT 방식처럼 연속 token을 가리는 N-gram masking도 적용했다.
기존 알고리즘: mask 가능한 sequence를 확률에 따라 N-gram으로 자르고 후보를 shuffle한 뒤 필요한 token 수만큼 선택했다. TensorFlow의 `while_loop`와 여러 op로 구현하자 이 단계가 전처리 시간의 절반 이상을 썼다.
적용 기준: Custom Op은 C++ 구현 비용이 있지만 원하는 연산을 하나의 효율적인 kernel로 만들 수 있다. 이 사례는 기존 op 조합의 제어 흐름과 중간 연산 비용이 큰 경우에 해당했다.
Metadata: `REGISTER_OP("GetMaskingPosition")`에 세 input, output, seed attribute와 문서를 선언했다. `SetShapeFn`은 input shape를 그대로 output shape로 전달했다.
Kernel 구현: `OpKernel`을 상속하고 `Compute`를 override했다. `context->input(index)`로 tensor reference를 받고 `flat<T>()` handle을 통해 값을 읽고 썼다.
Output 할당: `allocate_output`에 순번과 shape를 넘겨 Tensor pointer를 얻고, 같은 방식으로 flat handle을 만들어 계산 결과를 채웠다.
장치 등록: `REGISTER_KERNEL_BUILDER(Name(...).Device(DEVICE_CPU), ...)`로 선언과 구현을 연결했다. 데이터 전처리용이라 CPU kernel로 지정했고 GPU 연산이면 별도 device 구현이 필요하다.
빌드: `tf.sysconfig.get_compile_flags()`와 `get_link_flags()`로 현재 TensorFlow 설치와 맞는 option을 얻었다. g++에서 C++11, shared, fPIC, O2 옵션으로 `.so`를 생성했다.
Python 연동: `tf.load_op_library(path).get_masking_position`으로 함수를 가져왔다. Python 함수처럼 tensor와 `max_n`, `num_tokens_to_mask`를 전달했다.
벤치마크: sequence length 128, 최대 N-gram 3 조건에서 100회 실행했다. TensorFlow 구현은 총 0.4697초, Custom Op은 0.001714초였다.
효과: 문장 하나의 병목이 4.697ms에서 17.147μs로 줄었다. 가장 느렸던 masking 함수가 전체 변환 과정에서 가장 빠른 연산이 됐다.
Trade-off: Bazel 없이 유지보수하기 어렵고 op 이해와 debugging 시간이 커진다. 반면 model kernel fusion이나 CUDA backward까지 포함한 핵심 연산에는 DeepSpeed처럼 큰 최적화 여지가 있다.