-module(viva_tensor@quant@compression). -compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]). -define(FILEPATH, "src/viva_tensor/quant/compression.gleam"). -export([quantize_int8/1, quantize_q4/2, dequantize/1, create_memory_hierarchy/3, allocate_tensor/3, checkpoint_savings/3, create_streamed/2, load_chunk/2, unload_chunk/2, create_pool/0, pool_alloc/2, pool_free/2, demonstrate_compression/0, main/0]). -export_type([quant_format/0, compressed_tensor/0, tensor_location/0, memory_tier/0, memory_hierarchy/0, offload_policy/0, access_record/0, checkpoint/0, checkpoint_strategy/0, streamed_tensor/0, memory_pool/0]). -if(?OTP_RELEASE >= 27). -define(MODULEDOC(Str), -moduledoc(Str)). -define(DOC(Str), -doc(Str)). -else. -define(MODULEDOC(Str), -compile([])). -define(DOC(Str), -compile([])). -endif. ?MODULEDOC(false). -type quant_format() :: fp32 | fp16 | {int8, float()} | {quant4, integer(), list(float())} | {quant4_min, integer(), list(float()), list(float())}. -type compressed_tensor() :: {compressed_tensor, list(integer()), list(integer()), quant_format(), integer()}. -type tensor_location() :: {on_gpu, integer()} | on_ram | {on_disk, binary()} | {hybrid, float()}. -type memory_tier() :: {memory_tier, tensor_location(), float(), float(), float()}. -type memory_hierarchy() :: {memory_hierarchy, memory_tier(), memory_tier(), gleam@option:option(memory_tier()), float()}. -type offload_policy() :: keep_on_gpu | {offload_to_ram, float()} | {offload_to_disk, float(), binary()} | {smart_offload, list(access_record())}. -type access_record() :: {access_record, integer(), integer(), integer()}. -type checkpoint() :: {checkpoint, viva_tensor@tensor:tensor(), integer(), float()}. -type checkpoint_strategy() :: no_checkpoint | {every_n, integer()} | {large_layers_only, float()} | {adaptive, float()}. -type streamed_tensor() :: {streamed_tensor, integer(), list(integer()), list(integer()), list(integer()), integer(), quant_format()}. -type memory_pool() :: {memory_pool, list({integer(), integer()}), integer(), integer()}. -file("src/viva_tensor/quant/compression.gleam", 129). ?DOC(false). -spec get_shape(viva_tensor@tensor:tensor()) -> list(integer()). get_shape(T) -> case T of {tensor, _, Shape} -> Shape; {strided_tensor, _, Shape@1, _, _} -> Shape@1; {native_tensor, _, Shape@2} -> Shape@2 end. -file("src/viva_tensor/quant/compression.gleam", 742). ?DOC(false). -spec find_max_abs(list(float())) -> float(). find_max_abs(Data) -> gleam@list:fold( Data, +0.0, fun(Acc, V) -> Abs_v = gleam@float:absolute_value(V), case Abs_v > Acc of true -> Abs_v; false -> Acc end end ). -file("src/viva_tensor/quant/compression.gleam", 144). ?DOC(false). -spec quantize_int8(viva_tensor@tensor:tensor()) -> compressed_tensor(). quantize_int8(T) -> Data = viva_tensor@tensor:to_list(T), Shape = get_shape(T), Max_val = find_max_abs(Data), Scale = case Max_val > +0.0 of true -> case Max_val of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator -> 127.0 / Gleam@denominator end; false -> 1.0 end, Quantized = gleam@list:map( Data, fun(V) -> Scaled = V * Scale, Clamped = gleam@float:clamp(Scaled, -127.0, 127.0), erlang:round(Clamped) end ), Num_elements = gleam@list:fold(Shape, 1, fun(Acc, Dim) -> Acc * Dim end), Memory = Num_elements + 4, {compressed_tensor, Quantized, Shape, {int8, Scale}, Memory}. -file("src/viva_tensor/quant/compression.gleam", 184). ?DOC(false). -spec quantize_q4(viva_tensor@tensor:tensor(), integer()) -> compressed_tensor(). quantize_q4(T, Block_size) -> Data = viva_tensor@tensor:to_list(T), Shape = get_shape(T), Blocks = gleam@list:sized_chunk(Data, Block_size), {Quantized_blocks, Scales} = gleam@list:fold( Blocks, {[], []}, fun(Acc, Block) -> {Q_acc, S_acc} = Acc, Block_max = find_max_abs(Block), Scale = case Block_max > +0.0 of true -> case Block_max of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator -> 15.0 / Gleam@denominator end; false -> 1.0 end, Q_block = gleam@list:map( Block, fun(V) -> Scaled = (V * Scale) + 8.0, Clamped = gleam@float:clamp(Scaled, +0.0, 15.0), erlang:round(Clamped) end ), {lists:append(Q_acc, Q_block), [Scale | S_acc]} end ), Num_elements = gleam@list:fold(Shape, 1, fun(Acc@1, Dim) -> Acc@1 * Dim end), Num_blocks = case Block_size of 0 -> 0; Gleam@denominator@1 -> ((Num_elements + Block_size) - 1) div Gleam@denominator@1 end, Memory = (Num_elements div 2) + (Num_blocks * 4), {compressed_tensor, Quantized_blocks, Shape, {quant4, Block_size, lists:reverse(Scales)}, Memory}. -file("src/viva_tensor/quant/compression.gleam", 280). ?DOC(false). -spec create_tensor(list(float()), list(integer())) -> viva_tensor@tensor:tensor(). create_tensor(Data, Shape) -> {tensor, Data, Shape}. -file("src/viva_tensor/quant/compression.gleam", 284). ?DOC(false). -spec get_at_index(list(float()), integer(), float()) -> float(). get_at_index(Lst, Idx, Default) -> case gleam@list:drop(Lst, Idx) of [First | _] -> First; [] -> Default end. -file("src/viva_tensor/quant/compression.gleam", 236). ?DOC(false). -spec dequantize(compressed_tensor()) -> viva_tensor@tensor:tensor(). dequantize(Ct) -> case erlang:element(4, Ct) of fp32 -> create_tensor( gleam@list:map(erlang:element(2, Ct), fun erlang:float/1), erlang:element(3, Ct) ); fp16 -> create_tensor( gleam@list:map(erlang:element(2, Ct), fun erlang:float/1), erlang:element(3, Ct) ); {int8, Scale} -> Data = gleam@list:map(erlang:element(2, Ct), fun(Q) -> case Scale of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator -> erlang:float(Q) / Gleam@denominator end end), create_tensor(Data, erlang:element(3, Ct)); {quant4, Block_size, Scales} -> Blocks = gleam@list:sized_chunk(erlang:element(2, Ct), Block_size), Data@1 = begin _pipe = gleam@list:index_map( Blocks, fun(Block, Idx) -> Scale@1 = get_at_index(Scales, Idx, 1.0), gleam@list:map(Block, fun(Q@1) -> case Scale@1 of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator@1 -> (erlang:float(Q@1) - 8.0) / Gleam@denominator@1 end end) end ), lists:append(_pipe) end, create_tensor(Data@1, erlang:element(3, Ct)); {quant4_min, Block_size@1, Scales@1, Mins} -> Blocks@1 = gleam@list:sized_chunk( erlang:element(2, Ct), Block_size@1 ), Data@2 = begin _pipe@1 = gleam@list:index_map( Blocks@1, fun(Block@1, Idx@1) -> Scale@2 = get_at_index(Scales@1, Idx@1, 1.0), Min = get_at_index(Mins, Idx@1, +0.0), gleam@list:map(Block@1, fun(Q@2) -> (case Scale@2 of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator@2 -> erlang:float(Q@2) / Gleam@denominator@2 end) + Min end) end ), lists:append(_pipe@1) end, create_tensor(Data@2, erlang:element(3, Ct)) end. -file("src/viva_tensor/quant/compression.gleam", 296). ?DOC(false). -spec create_memory_hierarchy(float(), float(), gleam@option:option(binary())) -> memory_hierarchy(). create_memory_hierarchy(Vram_gb, Ram_gb, Disk_path) -> Gpu_tier = {memory_tier, {on_gpu, 0}, Vram_gb, +0.0, 1008.0}, Ram_tier = {memory_tier, on_ram, Ram_gb, +0.0, 51.2}, Disk_tier = case Disk_path of {some, Path} -> {some, {memory_tier, {on_disk, Path}, 1000.0, +0.0, 7.0}}; none -> none end, Effective = (Vram_gb * 4.0) + (Ram_gb * 4.0), {memory_hierarchy, Gpu_tier, Ram_tier, Disk_tier, Effective}. -file("src/viva_tensor/quant/compression.gleam", 345). ?DOC(false). -spec allocate_tensor(memory_hierarchy(), float(), offload_policy()) -> {tensor_location(), memory_hierarchy()}. allocate_tensor(Hierarchy, Tensor_size_gb, Policy) -> case Policy of keep_on_gpu -> Gpu_free = erlang:element(3, erlang:element(2, Hierarchy)) - erlang:element( 4, erlang:element(2, Hierarchy) ), case Tensor_size_gb =< Gpu_free of true -> New_gpu = begin _record = erlang:element(2, Hierarchy), {memory_tier, erlang:element(2, _record), erlang:element(3, _record), erlang:element(4, erlang:element(2, Hierarchy)) + Tensor_size_gb, erlang:element(5, _record)} end, {{on_gpu, 0}, {memory_hierarchy, New_gpu, erlang:element(3, Hierarchy), erlang:element(4, Hierarchy), erlang:element(5, Hierarchy)}}; false -> New_ram = begin _record@1 = erlang:element(3, Hierarchy), {memory_tier, erlang:element(2, _record@1), erlang:element(3, _record@1), erlang:element(4, erlang:element(3, Hierarchy)) + Tensor_size_gb, erlang:element(5, _record@1)} end, {on_ram, {memory_hierarchy, erlang:element(2, Hierarchy), New_ram, erlang:element(4, Hierarchy), erlang:element(5, Hierarchy)}} end; {offload_to_ram, Threshold} -> Gpu_usage = case erlang:element(3, erlang:element(2, Hierarchy)) of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator -> erlang:element( 4, erlang:element(2, Hierarchy) ) / Gleam@denominator end, case Gpu_usage < Threshold of true -> New_gpu@1 = begin _record@2 = erlang:element(2, Hierarchy), {memory_tier, erlang:element(2, _record@2), erlang:element(3, _record@2), erlang:element(4, erlang:element(2, Hierarchy)) + Tensor_size_gb, erlang:element(5, _record@2)} end, {{on_gpu, 0}, {memory_hierarchy, New_gpu@1, erlang:element(3, Hierarchy), erlang:element(4, Hierarchy), erlang:element(5, Hierarchy)}}; false -> New_ram@1 = begin _record@3 = erlang:element(3, Hierarchy), {memory_tier, erlang:element(2, _record@3), erlang:element(3, _record@3), erlang:element(4, erlang:element(3, Hierarchy)) + Tensor_size_gb, erlang:element(5, _record@3)} end, {on_ram, {memory_hierarchy, erlang:element(2, Hierarchy), New_ram@1, erlang:element(4, Hierarchy), erlang:element(5, Hierarchy)}} end; {offload_to_disk, Ram_threshold, Disk_path} -> Gpu_free@1 = erlang:element(3, erlang:element(2, Hierarchy)) - erlang:element( 4, erlang:element(2, Hierarchy) ), Ram_usage = case erlang:element(3, erlang:element(3, Hierarchy)) of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator@1 -> erlang:element( 4, erlang:element(3, Hierarchy) ) / Gleam@denominator@1 end, case Tensor_size_gb =< Gpu_free@1 of true -> New_gpu@2 = begin _record@4 = erlang:element(2, Hierarchy), {memory_tier, erlang:element(2, _record@4), erlang:element(3, _record@4), erlang:element(4, erlang:element(2, Hierarchy)) + Tensor_size_gb, erlang:element(5, _record@4)} end, {{on_gpu, 0}, {memory_hierarchy, New_gpu@2, erlang:element(3, Hierarchy), erlang:element(4, Hierarchy), erlang:element(5, Hierarchy)}}; false -> case Ram_usage < Ram_threshold of true -> New_ram@2 = begin _record@5 = erlang:element(3, Hierarchy), {memory_tier, erlang:element(2, _record@5), erlang:element(3, _record@5), erlang:element( 4, erlang:element(3, Hierarchy) ) + Tensor_size_gb, erlang:element(5, _record@5)} end, {on_ram, {memory_hierarchy, erlang:element(2, Hierarchy), New_ram@2, erlang:element(4, Hierarchy), erlang:element(5, Hierarchy)}}; false -> {{on_disk, Disk_path}, Hierarchy} end end; {smart_offload, _} -> {{on_gpu, 0}, Hierarchy} end. -file("src/viva_tensor/quant/compression.gleam", 465). ?DOC(false). -spec checkpoint_savings(integer(), float(), checkpoint_strategy()) -> float(). checkpoint_savings(Num_layers, Layer_size_mb, Strategy) -> Total_mb = erlang:float(Num_layers) * Layer_size_mb, case Strategy of no_checkpoint -> +0.0; {every_n, N} -> Checkpoint_pct = 1.0 - (case erlang:float(N) of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator -> 1.0 / Gleam@denominator end), Total_mb * Checkpoint_pct; {large_layers_only, Threshold} -> case Layer_size_mb > Threshold of true -> Total_mb * 0.7; false -> +0.0 end; {adaptive, Pressure} -> Total_mb * Pressure end. -file("src/viva_tensor/quant/compression.gleam", 511). ?DOC(false). -spec create_streamed(list(integer()), integer()) -> streamed_tensor(). create_streamed(Shape, Chunk_dim) -> Total_elements = gleam@list:fold(Shape, 1, fun(Acc, D) -> Acc * D end), Chunk_elements = Chunk_dim, Total_chunks = case Chunk_elements of 0 -> 0; Gleam@denominator -> ((Total_elements + Chunk_elements) - 1) div Gleam@denominator end, {streamed_tensor, erlang:unique_integer(), Shape, [Chunk_dim], [], Total_chunks, {int8, 1.0}}. -file("src/viva_tensor/quant/compression.gleam", 527). ?DOC(false). -spec load_chunk(streamed_tensor(), integer()) -> streamed_tensor(). load_chunk(St, Chunk_idx) -> case gleam@list:contains(erlang:element(5, St), Chunk_idx) of true -> St; false -> {streamed_tensor, erlang:element(2, St), erlang:element(3, St), erlang:element(4, St), [Chunk_idx | erlang:element(5, St)], erlang:element(6, St), erlang:element(7, St)} end. -file("src/viva_tensor/quant/compression.gleam", 536). ?DOC(false). -spec unload_chunk(streamed_tensor(), integer()) -> streamed_tensor(). unload_chunk(St, Chunk_idx) -> {streamed_tensor, erlang:element(2, St), erlang:element(3, St), erlang:element(4, St), gleam@list:filter(erlang:element(5, St), fun(C) -> C /= Chunk_idx end), erlang:element(6, St), erlang:element(7, St)}. -file("src/viva_tensor/quant/compression.gleam", 559). ?DOC(false). -spec create_pool() -> memory_pool(). create_pool() -> {memory_pool, [], 0, 0}. -file("src/viva_tensor/quant/compression.gleam", 565). ?DOC(false). -spec pool_alloc(memory_pool(), integer()) -> {memory_pool(), boolean()}. pool_alloc(Pool, Size) -> Found = gleam@list:find( erlang:element(2, Pool), fun(B) -> {S, Count} = B, (S =:= Size) andalso (Count > 0) end ), case Found of {ok, {S@1, _}} -> New_buffers = gleam@list:map( erlang:element(2, Pool), fun(B@1) -> {Bs, Bc} = B@1, case Bs =:= S@1 of true -> {Bs, Bc - 1}; false -> B@1 end end ), {{memory_pool, New_buffers, erlang:element(3, Pool) + 1, erlang:element(4, Pool)}, true}; {error, _} -> New_buffers@1 = [{Size, 0} | erlang:element(2, Pool)], {{memory_pool, New_buffers@1, erlang:element(3, Pool) + 1, erlang:element(4, Pool) + Size}, false} end. -file("src/viva_tensor/quant/compression.gleam", 608). ?DOC(false). -spec pool_free(memory_pool(), integer()) -> memory_pool(). pool_free(Pool, Size) -> New_buffers = case gleam@list:find( erlang:element(2, Pool), fun(B) -> {S, _} = B, S =:= Size end ) of {ok, _} -> gleam@list:map( erlang:element(2, Pool), fun(B@1) -> {Bs, Bc} = B@1, case Bs =:= Size of true -> {Bs, Bc + 1}; false -> B@1 end end ); {error, _} -> [{Size, 1} | erlang:element(2, Pool)] end, {memory_pool, New_buffers, erlang:element(3, Pool) - 1, erlang:element(4, Pool)}. -file("src/viva_tensor/quant/compression.gleam", 770). ?DOC(false). -spec float_to_string(float()) -> binary(). float_to_string(F) -> Rounded = erlang:float(erlang:round(F * 100.0)) / 100.0, gleam_stdlib:float_to_string(Rounded). -file("src/viva_tensor/quant/compression.gleam", 775). ?DOC(false). -spec location_to_string(tensor_location()) -> binary(). location_to_string(Loc) -> case Loc of {on_gpu, Id} -> <<"GPU #"/utf8, (erlang:integer_to_binary(Id))/binary>>; on_ram -> <<"RAM"/utf8>>; {on_disk, Path} -> <<<<"Disk("/utf8, Path/binary>>/binary, ")"/utf8>>; {hybrid, Pct} -> <<<<"Hybrid("/utf8, (float_to_string(Pct * 100.0))/binary>>/binary, "% GPU)"/utf8>> end. -file("src/viva_tensor/quant/compression.gleam", 752). ?DOC(false). -spec compute_quantization_error( viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor() ) -> float(). compute_quantization_error(Original, Restored) -> Orig_data = viva_tensor@tensor:to_list(Original), Rest_data = viva_tensor@tensor:to_list(Restored), {Sum_error, Count} = gleam@list:fold( gleam@list:zip(Orig_data, Rest_data), {+0.0, 0}, fun(Acc, Pair) -> {Sum, Cnt} = Acc, {O, R} = Pair, Error = gleam@float:absolute_value(O - R), {Sum + Error, Cnt + 1} end ), case Count > 0 of true -> case erlang:float(Count) of +0.0 -> +0.0; -0.0 -> -0.0; Gleam@denominator -> Sum_error / Gleam@denominator end; false -> +0.0 end. -file("src/viva_tensor/quant/compression.gleam", 640). ?DOC(false). -spec demonstrate_compression() -> nil. demonstrate_compression() -> gleam_stdlib:println( <<"====================================================================="/utf8>> ), gleam_stdlib:println(<<" INT8 QUANTIZATION - Jacob et al. (2017)"/utf8>>), gleam_stdlib:println( <<" 32-bit -> 8-bit = 4x compression with <0.5% accuracy loss"/utf8>> ), gleam_stdlib:println( <<"=====================================================================\n"/utf8>> ), Hierarchy = create_memory_hierarchy(24.0, 32.0, none), gleam_stdlib:println(<<"--- Hardware Configuration ---"/utf8>>), gleam_stdlib:println(<<" GPU: 24GB VRAM (RTX 4090, 1008 GB/s)"/utf8>>), gleam_stdlib:println(<<" RAM: 32GB DDR5 (51.2 GB/s)"/utf8>>), gleam_stdlib:println(<<" Physical: 56GB"/utf8>>), gleam_stdlib:println( <<<<" Effective with INT8: "/utf8, (float_to_string(erlang:element(5, Hierarchy)))/binary>>/binary, "GB"/utf8>> ), gleam_stdlib:println(<<""/utf8>>), gleam_stdlib:println(<<"--- Quantization Benchmark ---"/utf8>>), T = viva_tensor@tensor:random_uniform([1024, 512]), Original_size = (1024 * 512) * 4, Int8 = quantize_int8(T), Q4 = quantize_q4(T, 64), gleam_stdlib:println( <<<<" Original (FP32): "/utf8, (erlang:integer_to_binary(Original_size div 1024))/binary>>/binary, " KB"/utf8>> ), gleam_stdlib:println( <<<<" INT8 (4x): "/utf8, (erlang:integer_to_binary(erlang:element(5, Int8) div 1024))/binary>>/binary, " KB"/utf8>> ), gleam_stdlib:println( <<<<" Q4 block=64: "/utf8, (erlang:integer_to_binary(erlang:element(5, Q4) div 1024))/binary>>/binary, " KB"/utf8>> ), Restored = dequantize(Int8), Error = compute_quantization_error(T, Restored), gleam_stdlib:println( <<<<" INT8 error: "/utf8, (float_to_string(Error * 100.0))/binary>>/binary, "%"/utf8>> ), gleam_stdlib:println(<<"\n--- Memory Hierarchy ---"/utf8>>), gleam_stdlib:println(<<" Tier 1 (GPU): 1008 GB/s - hot tensors"/utf8>>), gleam_stdlib:println(<<" Tier 2 (RAM): 51.2 GB/s - warm tensors"/utf8>>), gleam_stdlib:println(<<" Tier 3 (Disk): 7 GB/s - cold tensors"/utf8>>), gleam_stdlib:println(<<"\n--- Allocation Example ---"/utf8>>), Policy = {offload_to_ram, 0.8}, {Loc1, H1} = allocate_tensor(Hierarchy, 10.0, Policy), gleam_stdlib:println( <<" 10GB tensor -> "/utf8, (location_to_string(Loc1))/binary>> ), {Loc2, H2} = allocate_tensor(H1, 10.0, Policy), gleam_stdlib:println( <<" 10GB tensor -> "/utf8, (location_to_string(Loc2))/binary>> ), {Loc3, _} = allocate_tensor(H2, 10.0, Policy), gleam_stdlib:println( <<" 10GB tensor -> "/utf8, (location_to_string(Loc3))/binary>> ), gleam_stdlib:println(<<"\n--- Gradient Checkpointing ---"/utf8>>), Layers = 24, Layer_mb = 100.0, Total_mb = erlang:float(Layers) * Layer_mb, Savings_n2 = checkpoint_savings(Layers, Layer_mb, {every_n, 2}), Savings_n4 = checkpoint_savings(Layers, Layer_mb, {every_n, 4}), gleam_stdlib:println( <<<<" 24 layers x 100MB = "/utf8, (float_to_string(Total_mb))/binary>>/binary, "MB"/utf8>> ), gleam_stdlib:println( <<<<" EveryN(2): saves "/utf8, (float_to_string(Savings_n2))/binary>>/binary, "MB (50% memory, 33% more compute)"/utf8>> ), gleam_stdlib:println( <<<<" EveryN(4): saves "/utf8, (float_to_string(Savings_n4))/binary>>/binary, "MB (75% memory, 25% more compute)"/utf8>> ), gleam_stdlib:println( <<"\n====================================================================="/utf8>> ), gleam_stdlib:println(<<" SUMMARY: 24GB VRAM + 32GB RAM with INT8"/utf8>>), gleam_stdlib:println(<<" - Physical: 56GB"/utf8>>), gleam_stdlib:println(<<" - Effective: 224GB (4x from INT8)"/utf8>>), gleam_stdlib:println( <<" - Can fit: ~110B parameters (224GB / 2 bytes per param)"/utf8>> ), gleam_stdlib:println( <<" - RTX 4090 INT8 Tensor Cores: 660 TOPS vs 330 TOPS FP16"/utf8>> ), gleam_stdlib:println( <<"====================================================================="/utf8>> ). -file("src/viva_tensor/quant/compression.gleam", 636). ?DOC(false). -spec main() -> nil. main() -> demonstrate_compression().