개빡치는 정밀도 줄이기 & LoRA

def calc_mantissa(abs_x, exponent, normal_mask, MANTISSA_BITS, EXPONENT_BIAS, generator=None):
    mantissa_scaled = torch.where(
        normal_mask,
        (abs_x / (2.0 ** (exponent - EXPONENT_BIAS)) - 1.0) * (2**MANTISSA_BITS),
        (abs_x / (2.0 ** (-EXPONENT_BIAS + 1 - MANTISSA_BITS)))
    )

    mantissa_scaled += torch.rand(mantissa_scaled.size(), dtype=mantissa_scaled.dtype, layout=mantissa_scaled.layout, device=mantissa_scaled.device, generator=generator)
    return mantissa_scaled.floor() / (2**MANTISSA_BITS)

def string_to_seed(data):
    crc = 0xFFFFFFFF
    for byte in data:
        if isinstance(byte, str):
            byte = ord(byte)
        crc ^= byte
        for _ in range(8):
            if crc & 1:
                crc = (crc >> 1) ^ 0xEDB88320
            else:
                crc >>= 1
    return crc ^ 0xFFFFFFFF

def manual_stochastic_round_to_float8(x, dtype, generator=None):
    if dtype == torch.float8_e4m3fn:
        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 4, 3, 7
    elif dtype == torch.float8_e5m2:
        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 5, 2, 15
    else:
        raise ValueError("Unsupported dtype")

    x = x.half()
    sign = torch.sign(x)
    abs_x = x.abs()
    sign = torch.where(abs_x == 0, 0, sign)

    # Combine exponent calculation and clamping
    exponent = torch.clamp(
        torch.floor(torch.log2(abs_x)) + EXPONENT_BIAS,
        0, 2**EXPONENT_BITS - 1
    )

    # Combine mantissa calculation and rounding
    normal_mask = ~(exponent == 0)

    abs_x[:] = calc_mantissa(abs_x, exponent, normal_mask, MANTISSA_BITS, EXPONENT_BIAS, generator=generator)

    sign *= torch.where(
        normal_mask,
        (2.0 ** (exponent - EXPONENT_BIAS)) * (1.0 + abs_x),
        (2.0 ** (-EXPONENT_BIAS + 1)) * abs_x
    )

    inf = torch.finfo(dtype)
    torch.clamp(sign, min=inf.min, max=inf.max, out=sign)
    return sign

def stochastic_rounding(value, dtype, seed=0):
    if dtype == torch.float32:
        return value.to(dtype=torch.float32)
    if dtype == torch.float16:
        return value.to(dtype=torch.float16)
    if dtype == torch.bfloat16:
        return value.to(dtype=torch.bfloat16)
    if dtype == torch.float8_e4m3fn or dtype == torch.float8_e5m2:
        generator = torch.Generator(device=value.device)
        generator.manual_seed(seed)
        output = torch.empty_like(value, dtype=dtype)
        num_slices = max(1, (value.numel() / (4096 * 4096)))
        slice_size = max(1, round(value.shape[0] / num_slices))
        for i in range(0, value.shape[0], slice_size):
            output[i:i+slice_size].copy_(manual_stochastic_round_to_float8(value[i:i+slice_size], dtype, generator=generator))
        return output

    return value.to(dtype=dtype)

def stochastic_rounding(value, dtype, seed=0):
    ...
    generator = torch.Generator(device=value.device)
    generator.manual_seed(seed)
    
    ...
    return output

output = torch.empty_like(value, dtype=dtype)
output._copy(~~)

exponent = ...      # 2GB (임시 메모리 1)
mantissa = ...      # 2GB (임시 메모리 2)
noise = torch.rand(...) # 2GB (임시 메모리 3)
result = ...        # 1GB (최종 결과 FP8)

def stochastic_rounding(value, dtype, seed=0):
    ...
    
    output = torch.empty_like(value, dtype=dtype)
    num_slices = max(1, (value.numel() / (4096 * 4096)))
    slice_size = max(1, round(value.shape[0] / num_slices))
    for i in range(0, value.shape[0], slice_size):
        output[i:i+slice_size].copy_(manual_stochastic_round_to_float8(value[i:i+slice_size], dtype, generator=generator))
    return output

def manual_stochastic_round_to_float8(x, dtype, generator=None):
    # 1. 설정값 세팅 (FP8 종류에 따라 비트 수와 바이어스 결정)
    if dtype == torch.float8_e4m3fn:
        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 4, 3, 7
    elif dtype == torch.float8_e5m2:
        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 5, 2, 15
    
    # 2. 전처리
    x = x.half() # FP16으로 변환 (계산 효율)
    sign = torch.sign(x) # 부호(+/-) 저장해둠
    abs_x = x.abs()
    
    # 3. 지수(Exponent) 계산
    # log2를 취하면 지수(2의 몇 승인지)가 나옴.
    # 바이어스(Bias)를 더해주고, 비트 범위(0 ~ 2^4-1)를 넘지 않게 clamp(가위질)함.
    exponent = torch.clamp(
        torch.floor(torch.log2(abs_x)) + EXPONENT_BIAS,
        0, 2**EXPONENT_BITS - 1
    )

    # 4. 정규수/비정규수 구분 (normal_mask)
    # 지수가 0이면 "너무 작은 숫자"라서 특별 취급해야 함.
    normal_mask = ~(exponent == 0)

    # 5. 가수(Mantissa) 계산 (위에서 분석한 함수 호출)
    # ============  여기서 확률적으로 반올림해 bias 최소화함 ================
    abs_x[:] = calc_mantissa(abs_x, exponent, normal_mask, ...)
    # ============  여기서 확률적으로 반올림해 bias 최소화함 ================

    # 6. 최종 재조립 (Reconstruction)
    # 공식: (-1)^sign * 2^(exponent-bias) * (1 + mantissa)
    sign *= torch.where(
        normal_mask,
        (2.0 ** (exponent - EXPONENT_BIAS)) * (1.0 + abs_x), # 정규수: 1.xxx 복구
        (2.0 ** (-EXPONENT_BIAS + 1)) * abs_x                # 비정규수: 0.xxx 복구
    )

    # 7. 안전장치 (Clamping)
    # 확률적 반올림을 하다 보면 값이 튀어서 FP8 범위를 넘을 수 있음.
    # FP8의 최소/최대값 안으로 강제로 욱여넣음.
    inf = torch.finfo(dtype)
    torch.clamp(sign, min=inf.min, max=inf.max, out=sign)
    return sign

def calc_mantissa(abs_x, exponent, normal_mask, MANTISSA_BITS, EXPONENT_BIAS, generator=None):
    # [수식 해부 1] 정규화 (Normalization)
    # 목표: 실수 x를 가수부 정수(0, 1, 2...)로 변환
    mantissa_scaled = torch.where(
        normal_mask,
        # Case A: 일반적인 숫자 (Normalized)
        # 공식: (값 / 2^지수 - 1) * 2^가수비트
        # 1. (abs_x / 2^exponent) -> 1.xxxx 형태가 됨
        # 2. (- 1.0) -> 0.xxxx (앞의 1 제거, 히든 비트 처리)
        # 3. (* 2^MANTISSA_BITS) -> 소수를 정수 범위로 뻥튀기 (예: 0~7)
        (abs_x / (2.0 ** (exponent - EXPONENT_BIAS)) - 1.0) * (2**MANTISSA_BITS),
        
        # Case B: 아주 작은 숫자 (Denormalized)
        # 0 근처의 숫자는 공식이 다름. 히든 비트 1이 없다고 가정.
        (abs_x / (2.0 ** (-EXPONENT_BIAS + 1 - MANTISSA_BITS)))
    )

    # [수식 해부 2] 노이즈 주입 (Stochastic Rounding의 핵심!)
    # 정수화된 실수(예: 3.7)에 0~1 사이의 난수를 더함.
    # 3.7 + 0.1(난수) = 3.8 -> floor -> 3 (내림)
    # 3.7 + 0.8(난수) = 4.5 -> floor -> 4 (올림)
    # 결과적으로 0.7의 확률로 4가 되고, 0.3의 확률로 3이 됨.
    mantissa_scaled += torch.rand(..., generator=generator)
    
    # [수식 해부 3] 복원
    # 정수로 만든 값을 다시 소수점 형태로 되돌림 (.floor()로 소수점 날림)
    return mantissa_scaled.floor() / (2**MANTISSA_BITS)