Sign In

FlashSampling: Fast and Memory-Efficient Exact Sampling

μž‘μ„±μž
  • Haebom
μΉ΄ν…Œκ³ λ¦¬
Empty

μ €μž

Tomas Ruiz, Zhen Qin, Yifan Zhang, Xuyang Shen, Yiran Zhong, Mengdi Wang

πŸ’‘ κ°œμš”

λŒ€κ·œλͺ¨ μ–΄νœ˜λ₯Ό λ‹€λ£¨λŠ” μ–Έμ–΄ λͺ¨λΈ λ””μ½”λ”©μ—μ„œ λ°œμƒν•˜λŠ” λ©”λͺ¨λ¦¬ μ˜€λ²„ν—€λ“œμ™€ 좔가적인 컀널 μ‹€ν–‰ 문제λ₯Ό ν•΄κ²°ν•˜κΈ° μœ„ν•΄, FlashSampling은 λ‘œμ§“(logits) ν…μ„œλ₯Ό HBM(High Bandwidth Memory)에 λͺ…μ‹œμ μœΌλ‘œ μ €μž₯ν•˜μ§€ μ•Šκ³  LM ν—€λ“œ ν–‰λ ¬ κ³±μ…ˆκ³Ό μƒ˜ν”Œλ§ 과정을 μœ΅ν•©ν•˜λŠ” μƒˆλ‘œμš΄ μƒ˜ν”Œλ§ 기법을 μ œμ•ˆν•©λ‹ˆλ‹€. 이 방법은 온칩(on-chip)μ—μ„œ νƒ€μΌλ³„λ‘œ λ‘œμ§“μ„ κ³„μ‚°ν•˜κ³ , Gumbel λ…Έμ΄μ¦ˆλ₯Ό λ”ν•˜λ©°, 각 ν–‰κ³Ό μ–΄νœ˜ 타일당 ν•˜λ‚˜μ˜ μ΅œλŒ€κ°’μ„ μœ μ§€ν•œ ν›„, 타일 κ°„μ˜ μž‘μ€ κ°μ†Œ 연산을 톡해 μ •ν™•ν•œ μƒ˜ν”Œλ§μ„ μˆ˜ν–‰ν•©λ‹ˆλ‹€.

πŸ”‘ μ‹œμ‚¬μ  및 ν•œκ³„

β€’
FlashSampling은 λ‘œμ§“ ν…μ„œλ₯Ό HBM에 μ €μž₯ν•  ν•„μš” 없이 LM ν—€λ“œ μ—°μ‚°κ³Ό μƒ˜ν”Œλ§μ„ ν†΅ν•©ν•˜μ—¬ λ©”λͺ¨λ¦¬ 접근을 크게 쀄이고 컀널 싀행을 κ°„μ†Œν™”ν•©λ‹ˆλ‹€.
β€’
ν…μ„œ 병렬 처리 ν™˜κ²½μ—μ„œ 기쑴의 λ‘œμ§“ μ˜¬κ²Œλ”(all-gather)λ₯Ό GPU κ°„ 슀트리밍 μ“°κΈ°(peer-to-peer writes)둜 λŒ€μ²΄ν•˜μ—¬, GPU κ°„ 톡신과 계산 및 HBM λ‘œλ“œλ₯Ό μ€‘μ²©μ‹œμΌœ μ„±λŠ₯을 크게 ν–₯μƒμ‹œν‚΅λ‹ˆλ‹€.
β€’
μ œμ•ˆλœ 기법은 근사 없이 μ •ν™•ν•œ μƒ˜ν”Œλ§μ„ μ œκ³΅ν•˜λ©°, λ‹€μ–‘ν•œ 데이터센터 GPUμ—μ„œ 컀널 μˆ˜μ€€μ˜ 속도 ν–₯상과 vLLM μ‹€ν—˜μ—μ„œ 좜λ ₯ 토큰당 μ‹œκ°„μ„ μ΅œλŒ€ 10%κΉŒμ§€ λ‹¨μΆ•ν•˜λŠ” μ„±κ³Όλ₯Ό λ³΄μž…λ‹ˆλ‹€.
β€’
ν˜„μž¬ μ—°κ΅¬λŠ” 주둜 ν…μ„œ 병렬 처리 ν™˜κ²½μ— μ΄ˆμ μ„ λ§žμΆ”κ³  있으며, κ·Έλ£Ήν™”λœ λ³€ν˜•μ€ 온라인 및 ν…μ„œ 병렬 μ„€μ •μ—μ„œ 정확성을 보μž₯ν•˜μ§€λ§Œ, 더 넓은 λ²”μœ„μ˜ λΆ„μ‚° ν™˜κ²½μ΄λ‚˜ λ³΅μž‘ν•œ λͺ¨λΈ ꡬ쑰에 λŒ€ν•œ 적용 κ°€λŠ₯μ„± 및 μ΅œμ ν™”λŠ” ν–₯ν›„ 연ꡬ 과제둜 λ‚¨μ•„μžˆμŠ΅λ‹ˆλ‹€.
πŸ‘
Made with Slashpage