/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/shardformer/layer/__init__.py
50 строк
1 KB
duanjunwen
[Feature] Support Distributed LogProb for GRPO Training (#6247)
05 авг 2025, 08:59
05 авг 2025, 08:59
4551853
Код
Авторство
О чём код?
from ._operation import all_to_all_comm from .attn import AttnMaskType, ColoAttention, RingAttention, get_pad_info from .dropout import DropoutForParallelInput, DropoutForReplicatedInput from .embedding import Embedding1D, PaddingEmbedding, VocabParallelEmbedding1D from .linear import Linear1D_Col, Linear1D_Row, LinearWithGradAccum, PaddingLMHead, VocabParallelLMHead1D from .loss import cross_entropy_1d, dist_cross_entropy, dist_log_prob, dist_log_prob_1d from .normalization import FusedLayerNorm, FusedRMSNorm, LayerNorm, RMSNorm from .parallel_module import ParallelModule from .qkv_fused_linear import ( FusedLinear, FusedLinear1D_Col, FusedLinear1D_Row, GPT2FusedLinearConv, GPT2FusedLinearConv1D_Col, GPT2FusedLinearConv1D_Row, ) __all__ = [ "Embedding1D", "VocabParallelEmbedding1D", "LinearWithGradAccum", "Linear1D_Col", "Linear1D_Row", "GPT2FusedLinearConv", "GPT2FusedLinearConv1D_Row", "GPT2FusedLinearConv1D_Col", "DropoutForParallelInput", "DropoutForReplicatedInput", "cross_entropy_1d", "dist_cross_entropy", "dist_log_prob_1d", "dist_log_prob", "BaseLayerNorm", "LayerNorm", "RMSNorm", "FusedLayerNorm", "FusedRMSNorm", "FusedLinear1D_Col", "FusedLinear", "ParallelModule", "PaddingEmbedding", "PaddingLMHead", "VocabParallelLMHead1D", "AttnMaskType", "ColoAttention", "RingAttention", "get_pad_info", "all_to_all_comm", "FusedLinear1D_Row", ]