Metadata-Version: 2.1
Name: torchastic
Version: 0.1.1
Summary: Stochastic bfloat16 based optimizer library.
Home-page: https://github.com/lodestone-rock/torchastic
Author: Lodestone
Author-email: lodestone.rock@gmail.com
License: UNKNOWN
Platform: UNKNOWN
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE

# torchastic
stochastic bfloat16 based optimizer library
## How to Use
```py
import torch
import torch.nn as nn
from torchastic import Compass, StochasticAccumulator


class Model(nn.Module):
    ...


# Init model
model = Model(*model_args)
optimizer = Compass(model.parameters(), lr=0.01, weight_decay=1e-2, amp_fac=5)

# Apply stochastic grad accumulator hooks
StochasticAccumulator.assign_hooks(model)

# Training
while True:

    # Gradient accumulation
    for _ in range(grad_accum_length):
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            loss = model.loss(*model_input)
        loss.backward()

    # Apply grad buffer back
    StochasticAccumulator.reassign_grad_buffer(model)
    optimizer.step()
    optimizer.zero_grad()
```

