Function all_to_all_single_fp8

colossalai/quantization/fp8.py:258–282

r""" This is wrapper for _all_to_all_single_fp8.

(
    output, input, output_split_sizes=None, input_split_sizes=None, fp8_format="e5m2", group=None, async_op=False
)

def all_to_all_single_fp8(
    output, input, output_split_sizes=None, input_split_sizes=None, fp8_format="e5m2", group=None, async_op=False
) -> Optional[Handle]:
    r"""
    This is wrapper for _all_to_all_single_fp8.
    """
    if process_group_is_intranode(group):
        return dist.all_to_all_single(
            output,
            input,
            output_split_sizes=output_split_sizes,
            input_split_sizes=input_split_sizes,
            group=group,
            async_op=async_op,
        )
    else:
        return _all_to_all_single_fp8(
            output,
            input,
            fp8_format=fp8_format,
            output_split_sizes=output_split_sizes,
            input_split_sizes=input_split_sizes,
            group=group,
            async_op=async_op,
        )


def cast_to_fp8_pipeline(inp: Any) -> None: