| Safe Haskell | None |
|---|---|
| Language | GHC2024 |
Circuit.LLM.Training
Description
Synopsis
- data AdamWState
- initAdamW :: GptConfig -> AdamWState
- adamwStep :: Double -> Double -> Double -> Double -> Double -> GptConfig -> AdamWState -> Gpt -> GptGrads -> (AdamWState, Gpt)
- updateVector :: Double -> Double -> Double -> Vector Double -> Vector Double -> Vector Double -> Vector Double
- updateMatrix :: Double -> Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
- trainStep :: GptConfig -> Gpt -> [Int] -> [Int] -> IO (Double, GptGrads)
- trainLoop :: Double -> Double -> Double -> Double -> Double -> GptConfig -> Gpt -> [Int] -> Int -> Int -> IO (Gpt, [Double])
- crossEntropyLoss :: Matrix Double -> [Int] -> Double
- trainStepMasked :: GptConfig -> Gpt -> [Int] -> [Int] -> [Bool] -> IO (Double, GptGrads)
- trainLoopMasked :: Double -> Double -> Double -> Double -> Double -> GptConfig -> Gpt -> [Int] -> Int -> Int -> Double -> Int -> IO (Gpt, [Double])
- maskPositions :: Int -> Double -> IO [Bool]
- applyMask :: [Int] -> [Bool] -> Int -> [Int]
Optimizer
data AdamWState Source #
AdamW optimizer state: step count, first moments, second moments.
initAdamW :: GptConfig -> AdamWState Source #
Initialize AdamW state.
adamwStep :: Double -> Double -> Double -> Double -> Double -> GptConfig -> AdamWState -> Gpt -> GptGrads -> (AdamWState, Gpt) Source #
One AdamW step: update model parameters given gradients.
Parameter updates (exported for metric-oracle tests)
updateVector :: Double -> Double -> Double -> Vector Double -> Vector Double -> Vector Double -> Vector Double Source #
AdamW update for one vector.
updateMatrix :: Double -> Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double Source #
AdamW update for one matrix: param -= lr_t * m / (sqrt(v) + eps) + wd_lr * param
Training
trainStep :: GptConfig -> Gpt -> [Int] -> [Int] -> IO (Double, GptGrads) Source #
One training step: forward + backward + loss, return (loss, gradients).
trainLoop :: Double -> Double -> Double -> Double -> Double -> GptConfig -> Gpt -> [Int] -> Int -> Int -> IO (Gpt, [Double]) Source #
Full training loop with AdamW optimization. data_ is a flat list of token IDs. We slide a window of seqLen tokens. Returns the trained model.
crossEntropyLoss :: Matrix Double -> [Int] -> Double Source #
AdamW update for one matrix. Training loop
Cross-entropy loss from logits and target token IDs.
Masked LM (BERT-style)
trainStepMasked :: GptConfig -> Gpt -> [Int] -> [Int] -> [Bool] -> IO (Double, GptGrads) Source #
One masked-LM training step: forward + backward + masked loss.
trainLoopMasked :: Double -> Double -> Double -> Double -> Double -> GptConfig -> Gpt -> [Int] -> Int -> Int -> Double -> Int -> IO (Gpt, [Double]) Source #
Masked-LM training loop with AdamW optimization.
- maskRate is the fraction of positions to mask per step.
- maskId is the token ID used to represent the [MASK] token.
- data_ is a flat list of token IDs. We slide a window of seqLen tokens.
- The targets are the original (unmasked) token IDs.