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

자세한 내용과 근거는 DAKER 원문에서 확인하실 수 있습니다.

핵심 요약

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

• 동일 입력·동일 라우팅 인덱스로 nested torch.cond 버전과 torch.switch 버전을 각각 만듭니다.

• 둘 다 torch.compile 또는 export 경로로 트레이스합니다.

원문 보기: https://daker.ai/community/pytorch-2-14-torch-switch-moe-graph-readability

DAKER: https://daker.ai/community/pytorch-2-14-torch-switch-moe-graph-readability