SkewAdam Mixture-of-Experts ਟ੍ਰੇਨਿੰਗ ਮੈਮੋਰੀ ਨੂੰ 60% ਤੋਂ ਵੱਧ ਘਟਾਉਂਦਾ ਹੈ ਅਤੇ ਸ਼ੁੱਧਤਾ (accuracy) ਵਿੱਚ ਇੱਕ ਮਹੱਤਵਪੂਰਨ ਵਾਧਾ ਦਿੰਦਾ ਹੈ, ਜਿਸ ਨਾਲ ਡਿਵੈਲਪਰ ਇੱਕ 6.78 ਬਿਲੀਅਨ-ਪੈਰਾਮੀਟਰ ਵਾਲੇ MoE ਮਾਡਲ ਨੂੰ ਇੱਕ ਸਿੰਗਲ 40 GB GPU 'ਤੇ ਚਲਾ ਸਕਦੇ ਹਨ।

MoE ਟ੍ਰੇਨਿੰਗ ਮੈਮੋਰੀ ਦੀਆਂ ਸੀਮਾਵਾਂ (memory wall) ਕਿਉਂ ਮਹਿਸੂਸ ਕਰਦੀ ਹੈ

Mixture-of-Experts ਆਰਕੀਟੈਕਚਰ ਇੱਕ ਡੈਂਸ ਬੈਕਬੋਨ (dense backbone) ਨੂੰ ਬਰਕਰਾਰ ਰੱਖਦੇ ਹੋਏ, ਹਰੇਕ ਇਨਪੁਟ ਨੂੰ "ਐਕਸਪਰਟ" (expert) ਸਬ-ਨੈੱਟਵਰਕਾਂ ਦੇ ਇੱਕ ਛੋਟੇ ਹਿੱਸੇ ਰਾਹੀਂ ਭੇਜਦੇ ਹਨ। ਇਸ ਦਾ ਫਾਇਦਾ ਇਹ ਹੈ ਕਿ ਇੱਕ ਅਜਿਹਾ ਮਾਡਲ ਮਿਲਦਾ ਹੈ ਜੋ ਕੰਪਿਊਟ (compute) ਵਿੱਚ ਅਨੁਪਾਤਕ ਵਾਧੇ ਤੋਂ ਬਿਨਾਂ ਅਰਬਾਂ ਪੈਰਾਮੀਟਰਾਂ ਤੱਕ ਵਧ ਸਕਦਾ ਹੈ। ਅਸਲ ਵਿੱਚ, ਆਪਟੀਮਾਈਜ਼ਰ—ਖਾਸ ਕਰਕੇ AdamW ਵੇਰੀਐਂਟ ਜਿਸਦੀ ਜ਼ਿਆਦਾਤਰ ਟੀਮਾਂ ਵਰਤਦੀਆਂ ਹਨ—GPU RAM ਦਾ ਵੱਡਾ ਹਿੱਸਾ ਵਰਤ ਲੈਂਦਾ ਹੈ। 6.78 B-ਪੈਰਾਮੀਟਰ ਵਾਲੇ ਮਾਡਲ ਲਈ, ਸਿਰਫ਼ ਆਪਟੀਮਾਈਜ਼ਰ ਦੇ ਮੋਮੈਂਟਮ (momentum) ਅਤੇ ਵੇਰੀਐਂਸ (variance) ਟੈਂਸਰ ਹੀ ਲਗਭਗ 50 GB ਦੀ ਮੰਗ ਕਰਦੇ ਹਨ। ਜੇਕਰ ਐਕਟੀਵੇਸ਼ਨਜ਼ (activations) ਅਤੇ ਮਾਡਲ ਵੇਟਸ (weights) ਨੂੰ ਜੋੜ ਦਿੱਤਾ ਜਾਵੇ, ਤਾਂ ਪੀਕ ਮੈਮੋਰੀ 81.4 GB ਤੱਕ ਪਹੁੰਚ ਜਾਂਦੀ ਹੈ, ਜਿਸ ਕਾਰਨ ਮਲਟੀ-GPU ਸੈੱਟਅੱਪ ਜਾਂ ਹੌਲੀ ਟ੍ਰੇਨਿੰਗ ਸ਼ਡਿਊਲ ਦੀ ਲੋੜ ਪੈਂਦੀ ਹੈ।

SkewAdam ਕੀ ਵੱਖਰਾ ਕਰਦਾ ਹੈ

SkewAdam ਕੋਈ ਨਵਾਂ ਲਰਨਿੰਗ ਨਿਯਮ ਨਹੀਂ ਬਣਾਉਂਦਾ। ਇਹ Adam ਦੇ ਮੋਮੈਂਟਸ (moments) ਦੇ ਸਥਾਨ ਨੂੰ ਮੁੜ ਵਿਵਸਥਿਤ ਕਰਦਾ ਹੈ:

  • ਡੈਂਸ ਬੈਕਬੋਨ ਫੁੱਲ-ਪ੍ਰੀਸੀਜ਼ਨ ਮੋਮੈਂਟਮ ਨੂੰ ਰੱਖਦਾ ਹੈ, ਜੋ ਕਿ ਡੈਂਸ ਲੇਅਰਾਂ ਲਈ ਲੋੜੀਂਦੇ ਸਮੂਥ ਅਪਡੇਟਸ ਨੂੰ ਬਰਕਰਾਰ ਰੱਖਦਾ ਹੈ।
  • ਐਕਸਪਰਟ ਬੈਂਕ ਵੇਰੀਐਂਸ ਦੇ ਇੱਕ ਫੈਕਟਰਡ ਅਪ੍ਰੌਕਸੀਮੇਸ਼ਨ (factored approximation) ਨੂੰ ਸਟੋਰ ਕਰਦਾ ਹੈ, ਜਿਸ ਨਾਲ ਹਰੇਕ ਐਕਸਪਰਟ ਲਈ ਰੱਖੇ ਗਏ ਡੇਟਾ ਵਿੱਚ ਕਟੌਤੀ ਹੁੰਦੀ ਹੈ।
  • ਰੂਟਰ, ਜੋ ਇਹ ਫੈਸਲਾ ਕਰਦਾ ਹੈ ਕਿ ਕਿਹੜੇ ਐਕਸਪਰਟਸ ਨੂੰ ਐਕਟੀਵੇਟ ਕਰਨਾ ਹੈ, ਇੱਕ ਸਹੀ ਸੈਕੰਡ-ਮੋਮੈਂਟ ਅਨੁਮਾਨ (second-moment estimate) ਰੱਖਦਾ ਹੈ।

ਇਹਨਾਂ ਤਿੰਨਾਂ ਹਿੱਸਿਆਂ ਨੂੰ ਵੱਖ-ਵੱਖ "ਟਾਇਰਾਂ" (tiers) ਵਜੋਂ ਮੰਨ ਕੇ, ਆਪਟੀਮਾਈਜ਼ਰ ਉੱਥੇ ਵਾਧੂ ਪ੍ਰੀਸੀਜ਼ਨ ਨੂੰ ਹਟਾ ਦਿੰਦਾ ਹੈ ਜਿੱਥੇ ਇਸਦੀ ਲੋੜ ਘੱਟ ਹੁੰਦੀ ਹੈ ਅਤੇ ਉੱਥੇ ਇਸਨੂੰ ਬਰਕਰਾਰ ਰੱਖਦਾ ਹੈ ਜਿੱਥੇ ਇਹ ਸਭ ਤੋਂ ਵੱਧ ਮਹੱਤਵਪੂਰਨ ਹੁੰਦਾ ਹੈ।

ਮਾਪਣਯੋਗ ਪ੍ਰਭਾਵ (Measurable impact)

SkewAdam ਦੇ ਨਾਲ ਉਹੀ 6.78 B-ਪੈਰਾਮੀਟਰ ਵਾਲਾ MoE ਮਾਡਲ ਚਲਾਉਣ ਨਾਲ ਇਹ ਨਤੀਜੇ ਮਿਲਦੇ ਹਨ:

  • ਪੀਕ GPU ਮੈਮੋਰੀ: 31.3 GB (81.4 GB ਤੋਂ ਘੱਟ)
  • ਆਪਟੀਮਾਈਜ਼ਰ-ਸਟੇਟ ਫੁੱਟਪ੍ਰਿੰਟ: 1.29 GB (50 GB ਤੋਂ ਘੱਟ)

ਘਟਾਈ ਗਈ ਸਟੇਟ ਇੱਕ ਸਿੰਗਲ 40 GB ਐਕਸਲਰੇਟਰ ਦੇ ਅੰਦਰ ਆਰਾਮ ਨਾਲ ਫਿੱਟ ਹੋ ਜਾਂਦੀ ਹੈ।

ਸਿਰਫ਼ ਬਚਤ ਹੀ ਨਹੀਂ, ਸ਼ੁੱਧਤਾ ਵਿੱਚ ਵਾਧਾ ਵੀ

ਮੈਮੋਰੀ ਵਿੱਚ ਕਟੌਤੀ ਅਕਸਰ ਮਾਡਲ ਦੀ ਗੁਣਵੱਤਾ ਨੂੰ ਨੁਕਸਾਨ ਪਹੁੰਚਾਉਂਦੀ ਹੈ, ਪਰ SkewAdam perplexity—ਜੋ ਕਿ ਇੱਕ ਮਿਆਰੀ ਭਾਸ਼ਾ-ਮਾਡਲ ਮੈਟ੍ਰਿਕ ਹੈ—ਨੂੰ 126.8 (AdamW) ਤੋਂ ਵਧਾ ਕੇ 108.4 ਕਰ ਦਿੰਦਾ ਹੈ। ਇਹ ਉਸੇ ਕੰਮ 'ਤੇ Muon (120.2) ਅਤੇ Lion (393.7) ਤੋਂ ਵੀ ਬਿਹਤਰ ਪ੍ਰਦਰਸ਼ਨ ਕਰਦਾ ਹੈ। ਲੇਖਕਾਂ ਦਾ ਕਹਿਣਾ ਹੈ ਕਿ ਇਹ ਵਾਧਾ ਤਿੰਨਾਂ ਟਾਇਰਾਂ ਵਿੱਚ ਮੋਮੈਂਟਮ ਨੂੰ ਬਰਕਰਾਰ ਰੱਖਣ ਕਾਰਨ ਹੁੰਦਾ ਹੈ; ਉਹ ਤਰੀਕੇ ਜੋ ਮੋਮੈਂਟਮ ਨੂੰ ਪੂਰੀ ਤਰ੍ਹਾਂ ਹਟਾ ਦਿੰਦੇ ਹਨ, ਜਿਵੇਂ ਕਿ Adafactor, ਪਿੱਛੇ ਰਹਿ ਜਾਂਦੇ ਹਨ।

ਖੁੱਲ੍ਹੇ ਸਵਾਲ

SkewAdam ਦੇ ਨਤੀਜੇ 6.78 B-ਪੈਰਾਮੀਟਰ ਵਾਲੇ ਮਾਡਲ 'ਤੇ ਦਿਖਾਏ ਗਏ ਹਨ। ਇਹ ਅਜੇ ਵੀ ਅਸਪਸ਼ਟ ਹੈ ਕਿ ਕੀ ਉਹੀ ਮੈਮੋਰੀ-ਸਟੇਟ ਅਨੁਪਾਤ ਕਈ ਗੁਣਾ ਵੱਡੇ ਮਾਡਲਾਂ ਲਈ ਜਾਂ ਭਾਸ਼ਾ ਮਾਡਲਿੰਗ ਤੋਂ ਪਰੇ ਦੇ ਕੰਮਾਂ ਲਈ ਵੀ ਲਾਗੂ ਹੁੰਦੇ ਹਨ।

ਕਿਸ ਚੀਜ਼ 'ਤੇ ਨਜ਼ਰ ਰੱਖਣੀ ਚਾਹੀਦੀ ਹੈ

  • ਵੱਡੇ MoE ਕੌਂਫਿਗਰੇਸ਼ਨਾਂ (ਦਸਾਂ ਬਿਲੀਅਨ ਪੈਰਾਮੀਟਰਾਂ) 'ਤੇ ਬੈਂਚਮਾਰਕਸ।
  • ਓਪਨ-ਸੋਰਸ ਫਰੇਮਵਰਕਾਂ ਦੁਆਰਾ ਅਪਣਾਉਣਾ ਅਤੇ ਪ੍ਰਸਿੱਧ ਟ੍ਰੇਨਿੰਗ ਸਕ੍ਰਿਪਟਾਂ ਵਿੱਚ ਸ਼ਾਮਲ ਹੋਣਾ।
  • ਲੁਕਵੇਂ ਸਮਝੌਤਿਆਂ (trade-offs) 'ਤੇ ਭਾਈਚਾਰਕ ਫੀਡਬੈਕ, ਜਿਵੇਂ ਕਿ ਵੱਖ-ਵੱਖ ਲਰਨਿੰਗ-ਰੇਟ ਸ਼ਡਿਊਲਾਂ ਦੇ ਅਧੀਨ ਕਨਵਰਜੈਂਸ ਸਪੀਡ ਜਾਂ ਸਥਿਰਤਾ।

ਸਰੋਤ: ਡਿਵੈਲਪਰ ਦੀ ਪੋਸਟ ਜਿਸ ਵਿੱਚ ਆਪਟੀਮਾਈਜ਼ਰ ਦੇ ਡਿਜ਼ਾਈਨ ਅਤੇ ਅਨੁਭਵੀ ਨਤੀਜਿਆਂ ਦਾ ਵੇਰਵਾ ਦਿੱਤਾ ਗਿਆ ਹੈ।