Class AllReduceParams

tensorrt_llm/functional.py:3899–3938

class AllReduceParams():

    def __init__(self,
                 strategy: AllReduceStrategy = AllReduceStrategy.AUTO,
                 fusion_op: AllReduceFusionOp = AllReduceFusionOp.NONE,
                 bias: Optional[Tensor] = None,
                 residual: Optional[Tensor] = None,
                 norm_weight: Optional[Tensor] = None,
                 scale: Optional[Tensor] = None,
                 norm_pre_residual_weight: Optional[Tensor] = None,
                 eps: float = 1e-06,
                 enable_allreduce: bool = True,
                 trigger_completion_at_end: bool = True):
        self.strategy = strategy
        self.fusion_op = fusion_op
        self.bias = bias
        self.residual = residual
        self.norm_weight = norm_weight
        self.scale = scale
        self.norm_pre_residual_weight = norm_pre_residual_weight
        self.eps = eps
        # For torch path only, has no effect on TRT path
        self.enable_allreduce = enable_allreduce
        self.trigger_completion_at_end = trigger_completion_at_end
        assert fusion_op == AllReduceFusionOp.NONE.value or (residual
                                                             is not None)

    def has_affine(self):
        return 1 if self.norm_weight is not None else 0

    def has_bias(self):
        return 1 if self.bias is not None else 0

    def has_scale(self):
        return 1 if self.scale is not None else 0

    def update_strategy(self):
        if self.strategy == AllReduceStrategy.AUTO and default_net(
        ).plugin_config.user_buffer:
            self.strategy = AllReduceStrategy.UB


class MoEAllReduceParams(AllReduceParams):