pruning_utils

Classes

GQAInitMode
MlpInitMode
LinearInitMode
HiddenSizeInitMode

Functions

resolve_pruning_mixin

Convert pruning_mixin argument to PruningMixIn instance(s).

class GQAInitMode

Bases: Enum

AverageKV = 'AverageKV'

CopyAsIs = 'CopyAsIs'

Degrouping = 'Degrouping'

FirstKV = 'FirstKV'

PruneKVHeads = 'PruneKVHeads'

RandomBlock = 'RandomBlock'

RandomKV = 'RandomKV'

class HiddenSizeInitMode

Bases: Enum

CopyAsIs = 'CopyAsIs'

PruneByChannelRanking = 'PruneByChannelRanking'

Random = 'Random'

Truncate = 'Truncate'

class LinearInitMode

Bases: Enum

FromTeacher = 'FromTeacher'

Random = 'Random'

class MlpInitMode

Bases: Enum

ConcatExpertsIntoDenseFFN = 'ConcatExpertsIntoDenseFFN'

CopyAsIs = 'CopyAsIs'

ExpertRemoval = 'ExpertRemoval'

MoEChannelPruning = 'MoEChannelPruning'

PruneByActivationsLog = 'PruneByActivationsLog'

Random = 'Random'

Truncate = 'Truncate'

resolve_pruning_mixin(pruning_mixin, descriptor)

Convert pruning_mixin argument to PruningMixIn instance(s).

Parameters:

• pruning_mixin – Can be a string identifier, PruningMixIn instance, or a list of any of those types.

• descriptor (Type[ModelDescriptor]) – ModelDescriptor class that provides the pruning_mixins() mapping.

Returns:

PruningMixIn or List[PruningMixIn] depending on input type.

Return type:

PruningMixIn | List[PruningMixIn]