PyTorch 2.14 torch.switch로 다중 분기하면 nested cond보다 MoE 라우팅 그래프가 읽기 쉬워집니다 | DAKER 커뮤니티

한 줄 요약
PyTorch 2.14의 torch.switch는 여러 갈래를 한 연산으로 고르는 multi-way 분기 HOP입니다. MoE나 인덱스 라우팅처럼 분기가 많아지면 nested torch.cond로 깊게 쌓이던 트레이스 그래프가 한눈에 읽히기 쉬워집니다.

전문가 라우팅을 짜다 보면, 조건마다 torch.cond를 겹쳐 두기 쉽습니다. 갈래가 둘일 때는 그래도 버틸 만하지만, 전문가가 넷·여덟로 늘어나면 그래프가 계단처럼 깊어집니다. 오늘은 분기 개수가 세 갈래 이상인지부터 세어 보십시오. 갈래가 늘어날수록 switch 한 번이 nested cond보다 그래프를 짧게 만듭니다.

재현해 볼 실험 조건

1. 동일 입력·동일 라우팅 인덱스로 nested torch.cond 버전과 torch.switch 버전을 각각 만듭니다.
2. 둘 다 torch.compile 또는 export 경로로 트레이스합니다.
3. 생성된 그래프에서 분기 노드 깊이와 가독성을 나란히 비교합니다.

2.14에서는 multi-way 분기가 전용 HOP로 들어가, MoE·인덱스 라우팅처럼 갈래가 많은 패턴에서 nested cond가 막히던 자리를 메웁니다. 동작이 바뀌는 실험보다, 같은 라우팅을 그래프가 어떻게 표현하는지가 핵심입니다. 참가자 입장에서는 “전문가 선택” 로직을 한 블록으로 남기는지가 디버깅 시간을 가릅니다.

실무에서 바로 볼 포인트

• 갈래가 둘뿐이면 nested cond로도 충분한 경우가 많습니다. switch는 갈래가 늘 때 이득이 큽니다.
• 라우팅 인덱스 dtype·범위가 어긋나면 그래프가 깨지기 전에 런타임에서 먼저 실패합니다. 인덱스 범위를 먼저 고정하십시오.
• 어텐션·임베딩 경로와 함께 쓸 때는 라우팅 블록만 따로 compile해 그래프를 비교하는 편이 원인 분리에 유리합니다.
• 학습 노트에는 PyTorch 버전과 switch 사용 여부를 같이 적어 두면, 나중에 그래프 비교가 수월합니다.

관련 DAKER 학습

• 임베딩·셀프 어텐션 미니 실습
• PyTorch 실전 고급 트랙

설명용 생성 이미지입니다. 그래프 형태는 로컬 트레이스 결과로 확인하십시오.