Spaces:
Paused
Paused
File size: 3,845 Bytes
6f75f79 d9a8f1f 6f75f79 06d2e6b 6f75f79 ec0eb2c d9a8f1f ec0eb2c 06d2e6b 6f75f79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """memory-native — finite-state counter synapses + reversible activations, in pure PyTorch.
A training method that attacks all four memory pools of training (params, optimizer state,
gradients, activations) with two independent levers, implemented on stock PyTorch (CPU/CUDA),
with no custom engine:
* CompactCounterLinear / RMSCounterLinear -- a ternary weight whose optimizer lives inside
a per-synapse finite-state automaton; the update is fused into backward (no FP master, no
Adam moments, no full gradient buffer).
* ReversibleCouplingBlock -- activations are recomputed in backward instead of stored.
See README for what is measured vs what still needs the Triton/CUDA packed kernel.
"""
from .counter import (
C_DEFAULT,
CompactCounterLinear,
RMSCounterLinear,
decode_state,
encode_state,
stochastic_round,
ternary_gradient_unbiased,
)
from .reversible import ReversibleCouplingBlock, ReversibleSequence, ReversibleSequential
from .packed import PackedRMSCounterLinear, pack_codes, unpack_codes
from .slowfast import SlowFastCounterLinear
from .triton_counter import HAS_TRITON, TritonCounterLinear, triton_decode_matmul, triton_grad_x
from .baselines import TernaryQATLinear, make_linear
from .models import CONFIGS, GPT, GPTConfig
from .reversible_gpt import ReversibleGPT
from .fused_qkv import CounterQKVLinear
from .memory_ffn import CounterMemoryFFN, CounterValueMemory
from .moe_ffn import CounterMoEFFN
from .mtp import MultiTokenHead, MTPGPT, mtp_loss, shift_targets
from .group_counter import GroupCounterLinear
from .stack_linear import StackCounterLinear
from .mod import MoDBlock
from .glm import MNGLM, ReversibleMNGLM, GLMBlock, GLMAttention, RMSNorm
from .int8_compute import (
fp8_correlation,
int4_correlation,
int8_correlation,
int8_correlation_presaved,
int8_forward_ternary,
int8_mm,
quantize_int4_cols,
quantize_int8_cols,
quantize_int8_rows,
)
from .memory import compare_training_peak, fmt_bytes, memory_report, peak_training_memory
from .actquant import effective_bits, quantize_codes, stochastic_quantize
from .budget import BudgetRow, format_budget, training_budget
from .optimizers import GaLoreAdamW, LoMo, available_optimizers, build_optimizer
from .data import get_batch, load_corpus, synthetic_corpus
__version__ = "0.1.0"
__all__ = [
"C_DEFAULT",
"CompactCounterLinear",
"RMSCounterLinear",
"SlowFastCounterLinear",
"PackedRMSCounterLinear",
"pack_codes",
"unpack_codes",
"TritonCounterLinear",
"triton_decode_matmul",
"triton_grad_x",
"HAS_TRITON",
"encode_state",
"decode_state",
"stochastic_round",
"ternary_gradient_unbiased",
"ReversibleCouplingBlock",
"ReversibleSequential",
"ReversibleSequence",
"TernaryQATLinear",
"make_linear",
"GPT",
"ReversibleGPT",
"CounterQKVLinear",
"CounterMemoryFFN",
"CounterValueMemory",
"CounterMoEFFN",
"MultiTokenHead",
"MTPGPT",
"mtp_loss",
"shift_targets",
"GroupCounterLinear",
"StackCounterLinear",
"MoDBlock",
"MNGLM",
"ReversibleMNGLM",
"GLMBlock",
"GLMAttention",
"RMSNorm",
"fp8_correlation",
"int4_correlation",
"int8_correlation",
"int8_correlation_presaved",
"int8_forward_ternary",
"int8_mm",
"quantize_int4_cols",
"quantize_int8_cols",
"quantize_int8_rows",
"GPTConfig",
"CONFIGS",
"memory_report",
"peak_training_memory",
"compare_training_peak",
"fmt_bytes",
"build_optimizer",
"available_optimizers",
"GaLoreAdamW",
"LoMo",
"stochastic_quantize",
"quantize_codes",
"effective_bits",
"training_budget",
"format_budget",
"BudgetRow",
"load_corpus",
"get_batch",
"synthetic_corpus",
"__version__",
]
|