{-# LANGUAGE OverloadedStrings #-}

-- | Training: loss, backprop, AdamW optimizer, and training loop.
module Circuit.LLM.Training
  ( -- * Optimizer
    AdamWState,
    initAdamW,
    adamwStep,

    -- * Parameter updates (exported for metric-oracle tests)
    updateVector,
    updateMatrix,

    -- * Training
    trainStep,
    trainLoop,
    crossEntropyLoss,

    -- * Masked LM (BERT-style)
    trainStepMasked,
    trainLoopMasked,
    maskPositions,
    applyMask,
  )
where

import Circuit.LLM.Backprop
  ( BlockGrads (..),
    GptGrads (..),
    addGptGrads,
    bertBackward,
    crossEntropyBwd,
    gptBackward,
    scaleGptGrads,
    zeroGptGrads,
  )
import Circuit.LLM.GPT
  ( FeedForward (..),
    Gpt (..),
    GptConfig (..),
    TransformerBlock (..),
  )
import Control.Monad (when)
import Debug.Trace (trace)
import Numeric.LinearAlgebra
  ( Matrix,
    Vector,
    cols,
    fromList,
    reshape,
    rows,
    toList,
  )
import Numeric.LinearAlgebra qualified as LA
import System.Random (newStdGen, randomRs)
import Text.Printf (printf)

----------------------------------------------------------------------
-- AdamW optimizer
----------------------------------------------------------------------

-- | AdamW optimizer state: step count, first moments, second moments.
data AdamWState = AdamWState
  { AdamWState -> Int
adamT :: !Int,
    AdamWState -> GptGrads
adamM :: !GptGrads,
    AdamWState -> GptGrads
adamV :: !GptGrads
  }

-- | Initialize AdamW state.
initAdamW :: GptConfig -> AdamWState
initAdamW :: GptConfig -> AdamWState
initAdamW GptConfig
cfg =
  AdamWState
    { adamT :: Int
adamT = Int
0,
      adamM :: GptGrads
adamM = GptConfig -> GptGrads
zeroGptGrads GptConfig
cfg,
      adamV :: GptGrads
adamV = GptConfig -> GptGrads
zeroGptGrads GptConfig
cfg
    }

-- | One AdamW step: update model parameters given gradients.
adamwStep ::
  Double ->
  Double ->
  Double ->
  Double ->
  Double ->
  GptConfig ->
  AdamWState ->
  Gpt ->
  GptGrads ->
  (AdamWState, Gpt)
adamwStep :: Double
-> Double
-> Double
-> Double
-> Double
-> GptConfig
-> AdamWState
-> Gpt
-> GptGrads
-> (AdamWState, Gpt)
adamwStep Double
lr Double
beta1 Double
beta2 Double
eps Double
wd GptConfig
cfg AdamWState
state Gpt
model GptGrads
grads =
  let t :: Int
t = AdamWState -> Int
adamT AdamWState
state Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      m :: GptGrads
m = [Char] -> GptGrads -> GptGrads
forall a. [Char] -> a -> a
trace [Char]
"momentUpdate..." (GptGrads -> GptGrads) -> GptGrads -> GptGrads
forall a b. (a -> b) -> a -> b
$ Double -> GptGrads -> GptGrads -> GptGrads
momentUpdate Double
beta1 (AdamWState -> GptGrads
adamM AdamWState
state) GptGrads
grads
      v :: GptGrads
v = [Char] -> GptGrads -> GptGrads
forall a. [Char] -> a -> a
trace [Char]
"momentUpdate2..." (GptGrads -> GptGrads) -> GptGrads -> GptGrads
forall a b. (a -> b) -> a -> b
$ Double -> GptGrads -> GptGrads -> GptGrads
momentUpdate2 Double
beta2 (AdamWState -> GptGrads
adamV AdamWState
state) GptGrads
grads
      lr_t :: Double
lr_t = Double
lr Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
beta2 Double -> Int -> Double
forall a b. (Fractional a, Integral b) => a -> b -> a
^^ Int
t) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
beta1 Double -> Int -> Double
forall a b. (Fractional a, Integral b) => a -> b -> a
^^ Int
t)
      model' :: Gpt
model' = [Char] -> Gpt -> Gpt
forall a. [Char] -> a -> a
trace [Char]
"applyUpdates..." (Gpt -> Gpt) -> Gpt -> Gpt
forall a b. (a -> b) -> a -> b
$ Double
-> Double
-> Double
-> GptConfig
-> Gpt
-> GptGrads
-> GptGrads
-> Gpt
applyUpdates Double
lr_t Double
eps (Double
wd Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
lr) GptConfig
cfg Gpt
model GptGrads
m GptGrads
v
      state' :: AdamWState
state' = AdamWState {adamT :: Int
adamT = Int
t, adamM :: GptGrads
adamM = GptGrads
m, adamV :: GptGrads
adamV = GptGrads
v}
   in (AdamWState
state', Gpt
model')

-- | Update first moment: m = beta1 * m + (1 - beta1) * grad.
momentUpdate :: Double -> GptGrads -> GptGrads -> GptGrads
momentUpdate :: Double -> GptGrads -> GptGrads -> GptGrads
momentUpdate Double
beta GptGrads
m GptGrads
g =
  GptGrads -> GptGrads -> GptGrads
addGptGrads (Double -> GptGrads -> GptGrads
scaleGptGrads Double
beta GptGrads
m) (Double -> GptGrads -> GptGrads
scaleGptGrads (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
beta) GptGrads
g)

-- | Update second moment: v = beta2 * v + (1 - beta2) * grad^2.
momentUpdate2 :: Double -> GptGrads -> GptGrads -> GptGrads
momentUpdate2 :: Double -> GptGrads -> GptGrads -> GptGrads
momentUpdate2 Double
beta GptGrads
v GptGrads
g =
  let sg :: GptGrads
sg = [Char] -> GptGrads -> GptGrads
forall a. [Char] -> a -> a
trace [Char]
"squareGrads..." (GptGrads -> GptGrads) -> GptGrads -> GptGrads
forall a b. (a -> b) -> a -> b
$ GptGrads -> GptGrads
squareGrads GptGrads
g
      scaled :: GptGrads
scaled = [Char] -> GptGrads -> GptGrads
forall a. [Char] -> a -> a
trace [Char]
"scaleGrads..." (GptGrads -> GptGrads) -> GptGrads -> GptGrads
forall a b. (a -> b) -> a -> b
$ Double -> GptGrads -> GptGrads
scaleGptGrads (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
beta) GptGrads
sg
   in [Char] -> GptGrads -> GptGrads
forall a. [Char] -> a -> a
trace [Char]
"momentUpdate2 add..." (GptGrads -> GptGrads) -> GptGrads -> GptGrads
forall a b. (a -> b) -> a -> b
$ GptGrads -> GptGrads -> GptGrads
addGptGrads (Double -> GptGrads -> GptGrads
scaleGptGrads Double
beta GptGrads
v) GptGrads
scaled

-- | Element-wise square of gradients (for second moment).
squareGrads :: GptGrads -> GptGrads
squareGrads :: GptGrads -> GptGrads
squareGrads GptGrads
g =
  let ok :: Bool
ok = Matrix Double -> Int
forall t. Matrix t -> Int
rows (GptGrads -> Matrix Double
ggWte GptGrads
g) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
   in if Bool -> Bool
not Bool
ok
        then [Char] -> GptGrads
forall a. HasCallStack => [Char] -> a
error [Char]
"squareGrads: empty Wte"
        else
          GptGrads
g
            { ggWte = ggWte g * ggWte g,
              ggWpe = ggWpe g * ggWpe g,
              ggBlocks = map squareBlockGrads (ggBlocks g),
              ggLnGamma = ggLnGamma g * ggLnGamma g,
              ggLnBeta = ggLnBeta g * ggLnBeta g,
              ggHead = ggHead g * ggHead g,
              ggHeadB = ggHeadB g * ggHeadB g
            }

squareBlockGrads :: BlockGrads -> BlockGrads
squareBlockGrads :: BlockGrads -> BlockGrads
squareBlockGrads BlockGrads
b =
  BlockGrads
b
    { bgAttnWq = bgAttnWq b * bgAttnWq b,
      bgAttnWk = bgAttnWk b * bgAttnWk b,
      bgAttnWv = bgAttnWv b * bgAttnWv b,
      bgAttnWo = bgAttnWo b * bgAttnWo b,
      bgAttnLnGamma = bgAttnLnGamma b * bgAttnLnGamma b,
      bgAttnLnBeta = bgAttnLnBeta b * bgAttnLnBeta b,
      bgFfnW1 = bgFfnW1 b * bgFfnW1 b,
      bgFfnB1 = bgFfnB1 b * bgFfnB1 b,
      bgFfnW2 = bgFfnW2 b * bgFfnW2 b,
      bgFfnB2 = bgFfnB2 b * bgFfnB2 b,
      bgFfnLnGamma = bgFfnLnGamma b * bgFfnLnGamma b,
      bgFfnLnBeta = bgFfnLnBeta b * bgFfnLnBeta b
    }

-- | Apply AdamW updates to all model parameters.
applyUpdates ::
  Double -> Double -> Double -> GptConfig -> Gpt -> GptGrads -> GptGrads -> Gpt
applyUpdates :: Double
-> Double
-> Double
-> GptConfig
-> Gpt
-> GptGrads
-> GptGrads
-> Gpt
applyUpdates Double
lr_t Double
eps Double
wd_lr GptConfig
_cfg Gpt
model GptGrads
m GptGrads
v =
  Gpt
model
    { gptWte = updateMatrix lr_t eps wd_lr (gptWte model) (ggWte m) (ggWte v),
      gptWpe = updateMatrix lr_t eps wd_lr (gptWpe model) (ggWpe m) (ggWpe v),
      gptBlocks =
        zipWith3
          (updateBlock lr_t eps wd_lr)
          (gptBlocks model)
          (ggBlocks m)
          (ggBlocks v),
      gptLnGamma = updateVector lr_t eps wd_lr (gptLnGamma model) (ggLnGamma m) (ggLnGamma v),
      gptLnBeta = updateVector lr_t eps wd_lr (gptLnBeta model) (ggLnBeta m) (ggLnBeta v),
      gptHead = updateMatrix lr_t eps wd_lr (gptHead model) (ggHead m) (ggHead v),
      gptHeadB = updateVector lr_t eps wd_lr (gptHeadB model) (ggHeadB m) (ggHeadB v)
    }

updateBlock :: Double -> Double -> Double -> TransformerBlock -> BlockGrads -> BlockGrads -> TransformerBlock
updateBlock :: Double
-> Double
-> Double
-> TransformerBlock
-> BlockGrads
-> BlockGrads
-> TransformerBlock
updateBlock Double
lr_t Double
eps Double
wd_lr TransformerBlock
tb BlockGrads
m BlockGrads
v =
  TransformerBlock
tb
    { tbAttnWq = updateMatrix lr_t eps wd_lr (tbAttnWq tb) (bgAttnWq m) (bgAttnWq v),
      tbAttnWk = updateMatrix lr_t eps wd_lr (tbAttnWk tb) (bgAttnWk m) (bgAttnWk v),
      tbAttnWv = updateMatrix lr_t eps wd_lr (tbAttnWv tb) (bgAttnWv m) (bgAttnWv v),
      tbAttnWo = updateMatrix lr_t eps wd_lr (tbAttnWo tb) (bgAttnWo m) (bgAttnWo v),
      tbAttnLnGamma = updateVector lr_t eps wd_lr (tbAttnLnGamma tb) (bgAttnLnGamma m) (bgAttnLnGamma v),
      tbAttnLnBeta = updateVector lr_t eps wd_lr (tbAttnLnBeta tb) (bgAttnLnBeta m) (bgAttnLnBeta v),
      tbFfn =
        let ff = TransformerBlock -> FeedForward
tbFfn TransformerBlock
tb; mf = BlockGrads
m; vf = BlockGrads
v
         in ff
              { ffW1 = updateMatrix lr_t eps wd_lr (ffW1 ff) (bgFfnW1 mf) (bgFfnW1 vf),
                ffB1 = updateVector lr_t eps wd_lr (ffB1 ff) (bgFfnB1 mf) (bgFfnB1 vf),
                ffW2 = updateMatrix lr_t eps wd_lr (ffW2 ff) (bgFfnW2 mf) (bgFfnW2 vf),
                ffB2 = updateVector lr_t eps wd_lr (ffB2 ff) (bgFfnB2 mf) (bgFfnB2 vf)
              },
      tbFfnLnGamma = updateVector lr_t eps wd_lr (tbFfnLnGamma tb) (bgFfnLnGamma m) (bgFfnLnGamma v),
      tbFfnLnBeta = updateVector lr_t eps wd_lr (tbFfnLnBeta tb) (bgFfnLnBeta m) (bgFfnLnBeta v)
    }

-- | AdamW update for one matrix: param -= lr_t * m / (sqrt(v) + eps) + wd_lr * param
updateMatrix :: Double -> Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
updateMatrix :: Double
-> Double
-> Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
updateMatrix Double
lr_t Double
eps Double
wd_lr Matrix Double
param Matrix Double
m_ Matrix Double
v_ =
  let nParams :: Int
nParams = Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
param Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
param
      nM :: Int
nM = Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
m_ Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
m_
      nV :: Int
nV = Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
v_ Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
v_
   in if Int
nParams Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nM Bool -> Bool -> Bool
|| Int
nParams Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nV
        then
          [Char] -> Matrix Double
forall a. HasCallStack => [Char] -> a
error ([Char] -> Matrix Double) -> [Char] -> Matrix Double
forall a b. (a -> b) -> a -> b
$
            [Char]
"updateMatrix: param="
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
param, Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
param)
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" m="
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
m_, Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
m_)
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" v="
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
v_, Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
v_)
        else
          let pv :: [Double]
pv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
param)
              mv :: [Double]
mv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
m_)
              vv :: [Double]
vv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
v_)
           in Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape
                (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
param)
                ( [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$
                    (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
p Double
m Double
v -> Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lr_t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
m Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double -> Double
forall a. Floating a => a -> a
sqrt Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
wd_lr Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p) [Double]
pv [Double]
mv [Double]
vv
                )

-- | AdamW update for one vector.
updateVector :: Double -> Double -> Double -> Vector Double -> Vector Double -> Vector Double -> Vector Double
updateVector :: Double
-> Double
-> Double
-> Vector Double
-> Vector Double
-> Vector Double
-> Vector Double
updateVector Double
lr_t Double
eps Double
wd_lr Vector Double
param Vector Double
m_ Vector Double
v_ =
  let pv :: [Double]
pv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
param
      mv :: [Double]
mv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
m_
      vv :: [Double]
vv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
v_
      np :: Int
np = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
pv
      nm :: Int
nm = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
mv
      nv :: Int
nv = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
vv
   in if Int
np Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nm Bool -> Bool -> Bool
|| Int
np Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nv
        then [Char] -> Vector Double
forall a. HasCallStack => [Char] -> a
error ([Char] -> Vector Double) -> [Char] -> Vector Double
forall a b. (a -> b) -> a -> b
$ [Char]
"updateVector size mismatch: param=" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
np [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" m=" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
nm [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" v=" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
nv
        else [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
p Double
m Double
v -> Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lr_t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
m Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double -> Double
forall a. Floating a => a -> a
sqrt Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
wd_lr Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p) [Double]
pv [Double]
mv [Double]
vv

-- | AdamW update for one matrix.
-- Training loop

----------------------------------------------------------------------

-- | Cross-entropy loss from logits and target token IDs.
crossEntropyLoss :: Matrix Double -> [Int] -> Double
crossEntropyLoss :: Matrix Double -> [Int] -> Double
crossEntropyLoss Matrix Double
logits [Int]
targetIds = (Double, Matrix Double) -> Double
forall a b. (a, b) -> a
fst (Matrix Double -> [Int] -> (Double, Matrix Double)
crossEntropyBwd Matrix Double
logits [Int]
targetIds)

-- | One training step: forward + backward + loss, return (loss, gradients).
trainStep :: GptConfig -> Gpt -> [Int] -> [Int] -> IO (Double, GptGrads)
trainStep :: GptConfig -> Gpt -> [Int] -> [Int] -> IO (Double, GptGrads)
trainStep GptConfig
cfg Gpt
model [Int]
inputIds [Int]
targetIds = do
  let (Double
loss, GptGrads
grads) = GptConfig -> Gpt -> [Int] -> [Int] -> (Double, GptGrads)
gptBackward GptConfig
cfg Gpt
model [Int]
inputIds [Int]
targetIds
  (Double, GptGrads) -> IO (Double, GptGrads)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
loss, GptGrads
grads)

-- | 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.
trainLoop ::
  Double ->
  Double ->
  Double ->
  Double ->
  Double ->
  GptConfig ->
  Gpt ->
  [Int] ->
  Int ->
  Int ->
  IO (Gpt, [Double])
trainLoop :: Double
-> Double
-> Double
-> Double
-> Double
-> GptConfig
-> Gpt
-> [Int]
-> Int
-> Int
-> IO (Gpt, [Double])
trainLoop Double
lr Double
beta1 Double
beta2 Double
eps Double
wd GptConfig
cfg Gpt
model [Int]
data_ Int
seqLen Int
steps =
  let opt0 :: AdamWState
opt0 = GptConfig -> AdamWState
initAdamW GptConfig
cfg
   in Gpt -> AdamWState -> [Double] -> Int -> Int -> IO (Gpt, [Double])
go Gpt
model AdamWState
opt0 [] Int
steps Int
0
  where
    go :: Gpt -> AdamWState -> [Double] -> Int -> Int -> IO (Gpt, [Double])
go Gpt
m AdamWState
_opt [Double]
losses Int
0 Int
_ = (Gpt, [Double]) -> IO (Gpt, [Double])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Gpt
m, [Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
losses)
    go Gpt
m AdamWState
opt [Double]
losses Int
n Int
offset = do
      let endIdx :: Int
endIdx = Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
seqLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      if Int
endIdx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
data_
        then Gpt -> AdamWState -> [Double] -> Int -> Int -> IO (Gpt, [Double])
go Gpt
m AdamWState
opt [Double]
losses (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
0
        else do
          let inputs :: [Int]
inputs = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
seqLen (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
offset [Int]
data_)
              targets :: [Int]
targets = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
seqLen (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
data_)
          (loss, grads) <- GptConfig -> Gpt -> [Int] -> [Int] -> IO (Double, GptGrads)
trainStep GptConfig
cfg Gpt
m [Int]
inputs [Int]
targets
          let (opt', m') = adamwStep lr beta1 beta2 eps wd cfg opt m grads
          let stepNum = Int
steps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          when (stepNum `mod` 10 == 0 || stepNum == 1) $
            printf "  step %d: loss=%.6f\n" stepNum loss
          go m' opt' (loss : losses) (n - 1) (offset + seqLen)

----------------------------------------------------------------------
-- Masked language model training (BERT-style)
----------------------------------------------------------------------

-- | One masked-LM training step: forward + backward + masked loss.
trainStepMasked :: GptConfig -> Gpt -> [Int] -> [Int] -> [Bool] -> IO (Double, GptGrads)
trainStepMasked :: GptConfig
-> Gpt -> [Int] -> [Int] -> [Bool] -> IO (Double, GptGrads)
trainStepMasked GptConfig
cfg Gpt
model [Int]
inputIds [Int]
targetIds [Bool]
mask = do
  let (Double
loss, GptGrads
grads) = GptConfig -> Gpt -> [Int] -> [Int] -> [Bool] -> (Double, GptGrads)
bertBackward GptConfig
cfg Gpt
model [Int]
inputIds [Int]
targetIds [Bool]
mask
  (Double, GptGrads) -> IO (Double, GptGrads)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
loss, GptGrads
grads)

-- | Randomly choose positions to mask.  At least one position is always
--   masked so the loss is well-defined.
maskPositions :: Int -> Double -> IO [Bool]
maskPositions :: Int -> Double -> IO [Bool]
maskPositions Int
seqLen Double
maskRate = do
  gen <- IO StdGen
forall (m :: * -> *). MonadIO m => m StdGen
newStdGen
  let nMask = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
seqLen Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
maskRate))
      idxs = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
nMask ((Int, Int) -> StdGen -> [Int]
forall g. RandomGen g => (Int, Int) -> g -> [Int]
forall a g. (Random a, RandomGen g) => (a, a) -> g -> [a]
randomRs (Int
0, Int
seqLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) StdGen
gen)
  pure [i `elem` idxs | i <- [0 .. seqLen - 1]]

-- | Replace masked positions with a fixed mask token ID.
applyMask :: [Int] -> [Bool] -> Int -> [Int]
applyMask :: [Int] -> [Bool] -> Int -> [Int]
applyMask [Int]
ids [Bool]
mask Int
maskId = (Int -> Bool -> Int) -> [Int] -> [Bool] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i Bool
m -> if Bool
m then Int
maskId else Int
i) [Int]
ids [Bool]
mask

-- | 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.
trainLoopMasked ::
  Double ->
  Double ->
  Double ->
  Double ->
  Double ->
  GptConfig ->
  Gpt ->
  [Int] ->
  Int ->
  Int ->
  Double ->
  Int ->
  IO (Gpt, [Double])
trainLoopMasked :: Double
-> Double
-> Double
-> Double
-> Double
-> GptConfig
-> Gpt
-> [Int]
-> Int
-> Int
-> Double
-> Int
-> IO (Gpt, [Double])
trainLoopMasked Double
lr Double
beta1 Double
beta2 Double
eps Double
wd GptConfig
cfg Gpt
model [Int]
data_ Int
seqLen Int
steps Double
maskRate Int
maskId =
  let opt0 :: AdamWState
opt0 = GptConfig -> AdamWState
initAdamW GptConfig
cfg
   in Gpt -> AdamWState -> [Double] -> Int -> Int -> IO (Gpt, [Double])
go Gpt
model AdamWState
opt0 [] Int
steps Int
0
  where
    go :: Gpt -> AdamWState -> [Double] -> Int -> Int -> IO (Gpt, [Double])
go Gpt
m AdamWState
_opt [Double]
losses Int
0 Int
_ = (Gpt, [Double]) -> IO (Gpt, [Double])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Gpt
m, [Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
losses)
    go Gpt
m AdamWState
opt [Double]
losses Int
n Int
offset = do
      if Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
seqLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
data_
        then Gpt -> AdamWState -> [Double] -> Int -> Int -> IO (Gpt, [Double])
go Gpt
m AdamWState
opt [Double]
losses (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
0
        else do
          let inputs :: [Int]
inputs = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
seqLen (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
offset [Int]
data_)
          mask <- Int -> Double -> IO [Bool]
maskPositions Int
seqLen Double
maskRate
          let maskedInputs = [Int] -> [Bool] -> Int -> [Int]
applyMask [Int]
inputs [Bool]
mask Int
maskId
          (loss, grads) <- trainStepMasked cfg m maskedInputs inputs mask
          let (opt', m') = adamwStep lr beta1 beta2 eps wd cfg opt m grads
          let stepNum = Int
steps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          when (stepNum `mod` 10 == 0 || stepNum == 1) $
            printf "  step %d: loss=%.6f\n" stepNum loss
          go m' opt' (loss : losses) (n - 1) (offset + seqLen)