Function copy_stochastic

optimizers/optimizer_utils.py:142–219

Performs stochastic rounding from source tensor to target tensor. Args: target: Destination tensor (determines the target format) source: Source tensor (typically float32) eps: Optional minimum value for stochastic rounding (for numerical stability)

(
    target: torch.Tensor,
    source: torch.Tensor,
    eps: Optional[float] = None
)

def copy_stochastic(
    target: torch.Tensor,
    source: torch.Tensor,
    eps: Optional[float] = None
) -> None:
    """
    Performs stochastic rounding from source tensor to target tensor.

    Args:
        target: Destination tensor (determines the target format)
        source: Source tensor (typically float32)
        eps: Optional minimum value for stochastic rounding (for numerical stability)
    """
    with torch.no_grad():
        # If target is float32, just copy directly
        if target.dtype == torch.float32:
            target.copy_(source)
            return

        # Special handling for int8
        if target.dtype == torch.int8:
            # Scale the source values to utilize the full int8 range
            scaled = source * 127.0  # Scale to [-127, 127]

            # Add random noise for stochastic rounding
            noise = torch.rand_like(scaled) - 0.5
            rounded = torch.round(scaled + noise)

            # Clamp to int8 range
            clamped = torch.clamp(rounded, -127, 127)
            target.copy_(clamped.to(torch.int8))
            return

        mantissa_bits, _ = get_format_params(target.dtype)

        # Convert source to int32 view
        source_int = source.view(dtype=torch.int32)

        # Calculate number of bits to round
        bits_to_round = 23 - mantissa_bits  # 23 is float32 mantissa bits

        # Create random integers for stochastic rounding
        rand = torch.randint_like(
            source,
            dtype=torch.int32,
            low=0,
            high=(1 << bits_to_round),
        )

        # Add random values to the bits that will be rounded off
        result = source_int.clone()
        result.add_(rand)

        # Mask to keep only the bits we want
        # Create mask with 1s in positions we want to keep
        mask = (-1) << bits_to_round
        result.bitwise_and_(mask)