@seanhogg/builderforce-memory-engine 2026.7.12 → 2026.7.14

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.d.ts CHANGED
@@ -10,9 +10,13 @@ export { MambaBlock } from './model/mamba1_block.js';
10
10
  export { SharedExpertMoE, LoadBalanceAccumulator, DEFAULT_MOE_CONFIG, DEFAULT_MOE_SEED, MoETrainer, EvermindModelPackage, } from './moe/index.js';
11
11
  export type { MoEConfig, MoEParam, RouteResult, MoESample, MoETrainOptions, MoEEpochResult, EvermindModelManifest, EvermindModelCard, EvermindModelType, EvermindModality, PackageMeta, ValidationResult, } from './moe/index.js';
12
12
  export { EvermindLM, EvermindLMTrainer, DEFAULT_LM_CONFIG, DEFAULT_LM_SEED } from './lm/index.js';
13
- export type { EvermindLMConfig, LMGenerateOptions, TextCodec } from './lm/index.js';
13
+ export type { EvermindLMConfig, LMGenerateOptions, TextCodec, EvermindLMTrainOptions } from './lm/index.js';
14
14
  export { AdamW } from './optim/adamw.js';
15
- export type { AdamWOptions, OptimTarget, OptimParam } from './optim/adamw.js';
15
+ export type { AdamWOptions, OptimTarget, OptimParam, ShardSpec } from './optim/adamw.js';
16
+ export { LoRAAdapter, EvermindLMLoRA, quantizeBase } from './training/lora.js';
17
+ export type { LoRAConfig, LoRAFitOptions, BaseQuant } from './training/lora.js';
18
+ export { DynamicLossScaler, roundFp16, fp16View } from './training/mixed_precision.js';
19
+ export type { LossScalerOptions } from './training/mixed_precision.js';
16
20
  export { MambaTrainer } from './training/trainer.js';
17
21
  export { Tensor, backward, enableGrad, noGrad, clearTape, recordOperation, crossEntropyLoss, crossEntropyGrad, } from './training/autograd.js';
18
22
  export { BPETokenizer } from './tokenizer/bpe.js';
@@ -1 +1 @@
1
- {"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;GAEG;AAIH,OAAO,EAAE,gBAAgB,EAAE,UAAU,EAAE,MAAM,wBAAwB,CAAC;AAGtE,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,cAAc,EAAE,MAAM,4BAA4B,CAAC;AAG5D,OAAO,EAAE,UAAU,EAAE,MAAM,yBAAyB,CAAC;AAGrD,OAAO,EACH,eAAe,EACf,sBAAsB,EACtB,kBAAkB,EAClB,gBAAgB,EAChB,UAAU,EACV,oBAAoB,GACvB,MAAM,gBAAgB,CAAC;AACxB,YAAY,EACR,SAAS,EACT,QAAQ,EACR,WAAW,EACX,SAAS,EACT,eAAe,EACf,cAAc,EACd,qBAAqB,EACrB,iBAAiB,EACjB,iBAAiB,EACjB,gBAAgB,EAChB,WAAW,EACX,gBAAgB,GACnB,MAAM,gBAAgB,CAAC;AAGxB,OAAO,EAAE,UAAU,EAAE,iBAAiB,EAAE,iBAAiB,EAAE,eAAe,EAAE,MAAM,eAAe,CAAC;AAClG,YAAY,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,SAAS,EAAE,MAAM,eAAe,CAAC;AACpF,OAAO,EAAE,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACzC,YAAY,EAAE,YAAY,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,kBAAkB,CAAC;AAI9E,OAAO,EAAE,YAAY,EAAE,MAAM,uBAAuB,CAAC;AACrD,OAAO,EACH,MAAM,EACN,QAAQ,EACR,UAAU,EACV,MAAM,EACN,SAAS,EACT,eAAe,EACf,gBAAgB,EAChB,gBAAgB,GACnB,MAAM,wBAAwB,CAAC;AAIhC,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAClD,YAAY,EAAE,gBAAgB,EAAE,OAAO,EAAE,wBAAwB,EAAE,qBAAqB,EAAE,MAAM,oBAAoB,CAAC;AAIrH,OAAO,EACH,eAAe,EACf,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,kBAAkB,EAClB,aAAa,EACb,aAAa,GAChB,MAAM,kBAAkB,CAAC;AAC1B,YAAY,EACR,qBAAqB,EACrB,SAAS,EACT,cAAc,EACd,KAAK,EACL,KAAK,GACR,MAAM,kBAAkB,CAAC;AAI1B,OAAO,EACH,cAAc,EACd,cAAc,EACd,iBAAiB,EACjB,oBAAoB,EACpB,UAAU,EACV,UAAU,EACV,UAAU,EACV,oBAAoB,EACpB,aAAa,EACb,iBAAiB,EACjB,YAAY,EACZ,kBAAkB,EAClB,MAAM,EACN,UAAU,GACb,MAAM,mBAAmB,CAAC;AAC3B,YAAY,EACR,YAAY,EACZ,UAAU,EACV,YAAY,EACZ,aAAa,EACb,MAAM,EACN,WAAW,EACX,YAAY,EACZ,UAAU,GACb,MAAM,mBAAmB,CAAC;AAI3B,OAAO,EACH,oBAAoB,EACpB,cAAc,EACd,qBAAqB,EACrB,oBAAoB,GACvB,MAAM,mBAAmB,CAAC;AAC3B,YAAY,EAAE,aAAa,EAAE,MAAM,mBAAmB,CAAC;AAIvD,OAAO,EACH,cAAc,EACd,mBAAmB,EACnB,aAAa,EACb,aAAa,EACb,cAAc,EACd,iBAAiB,EACjB,iBAAiB,EACjB,MAAM,IAAI,WAAW,EACrB,WAAW,EACX,UAAU,EACV,YAAY,EACZ,GAAG,GACN,MAAM,kBAAkB,CAAC;AAC1B,YAAY,EACR,WAAW,EACX,gBAAgB,EAChB,gBAAgB,EAChB,eAAe,EACf,gBAAgB,EAChB,wBAAwB,EACxB,uBAAuB,GAC1B,MAAM,kBAAkB,CAAC;AAI1B,OAAO,EAAE,KAAK,EAAE,gBAAgB,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MAAM,kBAAkB,CAAC;AAChG,YAAY,EAAE,QAAQ,EAAE,MAAM,kBAAkB,CAAC;AACjD,OAAO,EACH,eAAe,EACf,aAAa,EACb,iBAAiB,EACjB,mBAAmB,EACnB,eAAe,EACf,mBAAmB,GACtB,MAAM,kBAAkB,CAAC;AAC1B,YAAY,EAAE,QAAQ,EAAE,MAAM,kBAAkB,CAAC;AAIjD,OAAO,EAAE,SAAS,EAAE,WAAW,EAAE,KAAK,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AAI9E,YAAY,EACR,sBAAsB,EACtB,gBAAgB,EAChB,kBAAkB,EAClB,eAAe,EACf,SAAS,GACZ,MAAM,wBAAwB,CAAC;AAEhC,YAAY,EAAE,aAAa,EAAE,UAAU,EAAE,SAAS,EAAE,kBAAkB,EAAE,MAAM,2BAA2B,CAAC;AAC1G,YAAY,EAAE,iBAAiB,EAAE,UAAU,EAAE,UAAU,EAAE,kBAAkB,EAAE,gBAAgB,EAAE,MAAM,yBAAyB,CAAC;AAC/H,YAAY,EAAE,iBAAiB,EAAE,WAAW,EAAE,MAAM,yBAAyB,CAAC;AAC9E,YAAY,EAAE,iBAAiB,EAAE,WAAW,EAAE,MAAO,yBAAyB,CAAC;AAC/E,YAAY,EAAE,oBAAoB,EAAE,cAAc,EAAE,MAAM,4BAA4B,CAAC;AAIvF,OAAO,EACH,UAAU,EACV,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,qBAAqB,EACrB,eAAe,EACf,cAAc,EACd,UAAU,EACV,YAAY,EACZ,IAAI,EACJ,UAAU,GACb,MAAM,sBAAsB,CAAC;AAI9B,OAAO,EACH,YAAY,EACZ,cAAc,EACd,WAAW,EACX,WAAW,EACX,YAAY,EACZ,cAAc,EACd,sBAAsB,EACtB,wBAAwB,EACxB,cAAc,GACjB,MAAM,yBAAyB,CAAC;AAKjC,OAAO,EAAE,2BAA2B,EAAE,4BAA4B,EAAE,MAC3D,6BAA6B,CAAC;AACvC,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,qBAAqB,CAAC;AAC/B,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,gCAAgC,CAAC;AAC1C,OAAO,EAAE,kBAAkB,EAAE,cAAc,EAAE,MACpC,4BAA4B,CAAC;AACtC,OAAO,EAAE,gBAAgB,EAAE,yBAAyB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,MAC1F,0BAA0B,CAAC;AAGpC,OAAO,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MACrC,kBAAkB,CAAC;AAG5B,OAAO,EAAE,wBAAwB,EAAE,yBAAyB,EAAE,MACrD,0BAA0B,CAAC;AAGpC,OAAO,EAAE,sBAAsB,EAAE,uBAAuB,EAAE,YAAY,EAAE,MAC/D,wBAAwB,CAAC;AAIlC,OAAO,EACH,MAAM,EACN,UAAU,EACV,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,QAAQ,EACR,UAAU,EACV,YAAY,EACZ,aAAa,EACb,aAAa,EACb,mBAAmB,EACnB,WAAW,EACX,qBAAqB,EACrB,mBAAmB,EACnB,aAAa,GAChB,MAAM,mBAAmB,CAAC;AAC3B,YAAY,EACR,MAAM,EACN,aAAa,EACb,iBAAiB,EACjB,iBAAiB,EACjB,aAAa,EACb,WAAW,EACX,YAAY,EACZ,kBAAkB,GACrB,MAAM,mBAAmB,CAAC;AAE3B,OAAO,EAAE,kBAAkB,EAAE,MAAM,4BAA4B,CAAC;AAIhE,eAAO,MAAM,OAAO,UAAc,CAAC;AACnC,eAAO,MAAM,WAAW,wFAAwF,CAAC"}
1
+ {"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;GAEG;AAIH,OAAO,EAAE,gBAAgB,EAAE,UAAU,EAAE,MAAM,wBAAwB,CAAC;AAGtE,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,cAAc,EAAE,MAAM,4BAA4B,CAAC;AAG5D,OAAO,EAAE,UAAU,EAAE,MAAM,yBAAyB,CAAC;AAGrD,OAAO,EACH,eAAe,EACf,sBAAsB,EACtB,kBAAkB,EAClB,gBAAgB,EAChB,UAAU,EACV,oBAAoB,GACvB,MAAM,gBAAgB,CAAC;AACxB,YAAY,EACR,SAAS,EACT,QAAQ,EACR,WAAW,EACX,SAAS,EACT,eAAe,EACf,cAAc,EACd,qBAAqB,EACrB,iBAAiB,EACjB,iBAAiB,EACjB,gBAAgB,EAChB,WAAW,EACX,gBAAgB,GACnB,MAAM,gBAAgB,CAAC;AAGxB,OAAO,EAAE,UAAU,EAAE,iBAAiB,EAAE,iBAAiB,EAAE,eAAe,EAAE,MAAM,eAAe,CAAC;AAClG,YAAY,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,SAAS,EAAE,sBAAsB,EAAE,MAAM,eAAe,CAAC;AAC5G,OAAO,EAAE,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACzC,YAAY,EAAE,YAAY,EAAE,WAAW,EAAE,UAAU,EAAE,SAAS,EAAE,MAAM,kBAAkB,CAAC;AAGzF,OAAO,EAAE,WAAW,EAAE,cAAc,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAC/E,YAAY,EAAE,UAAU,EAAE,cAAc,EAAE,SAAS,EAAE,MAAM,oBAAoB,CAAC;AAChF,OAAO,EAAE,iBAAiB,EAAE,SAAS,EAAE,QAAQ,EAAE,MAAM,+BAA+B,CAAC;AACvF,YAAY,EAAE,iBAAiB,EAAE,MAAM,+BAA+B,CAAC;AAIvE,OAAO,EAAE,YAAY,EAAE,MAAM,uBAAuB,CAAC;AACrD,OAAO,EACH,MAAM,EACN,QAAQ,EACR,UAAU,EACV,MAAM,EACN,SAAS,EACT,eAAe,EACf,gBAAgB,EAChB,gBAAgB,GACnB,MAAM,wBAAwB,CAAC;AAIhC,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAClD,YAAY,EAAE,gBAAgB,EAAE,OAAO,EAAE,wBAAwB,EAAE,qBAAqB,EAAE,MAAM,oBAAoB,CAAC;AAIrH,OAAO,EACH,eAAe,EACf,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,kBAAkB,EAClB,aAAa,EACb,aAAa,GAChB,MAAM,kBAAkB,CAAC;AAC1B,YAAY,EACR,qBAAqB,EACrB,SAAS,EACT,cAAc,EACd,KAAK,EACL,KAAK,GACR,MAAM,kBAAkB,CAAC;AAI1B,OAAO,EACH,cAAc,EACd,cAAc,EACd,iBAAiB,EACjB,oBAAoB,EACpB,UAAU,EACV,UAAU,EACV,UAAU,EACV,oBAAoB,EACpB,aAAa,EACb,iBAAiB,EACjB,YAAY,EACZ,kBAAkB,EAClB,MAAM,EACN,UAAU,GACb,MAAM,mBAAmB,CAAC;AAC3B,YAAY,EACR,YAAY,EACZ,UAAU,EACV,YAAY,EACZ,aAAa,EACb,MAAM,EACN,WAAW,EACX,YAAY,EACZ,UAAU,GACb,MAAM,mBAAmB,CAAC;AAI3B,OAAO,EACH,oBAAoB,EACpB,cAAc,EACd,qBAAqB,EACrB,oBAAoB,GACvB,MAAM,mBAAmB,CAAC;AAC3B,YAAY,EAAE,aAAa,EAAE,MAAM,mBAAmB,CAAC;AAIvD,OAAO,EACH,cAAc,EACd,mBAAmB,EACnB,aAAa,EACb,aAAa,EACb,cAAc,EACd,iBAAiB,EACjB,iBAAiB,EACjB,MAAM,IAAI,WAAW,EACrB,WAAW,EACX,UAAU,EACV,YAAY,EACZ,GAAG,GACN,MAAM,kBAAkB,CAAC;AAC1B,YAAY,EACR,WAAW,EACX,gBAAgB,EAChB,gBAAgB,EAChB,eAAe,EACf,gBAAgB,EAChB,wBAAwB,EACxB,uBAAuB,GAC1B,MAAM,kBAAkB,CAAC;AAI1B,OAAO,EAAE,KAAK,EAAE,gBAAgB,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MAAM,kBAAkB,CAAC;AAChG,YAAY,EAAE,QAAQ,EAAE,MAAM,kBAAkB,CAAC;AACjD,OAAO,EACH,eAAe,EACf,aAAa,EACb,iBAAiB,EACjB,mBAAmB,EACnB,eAAe,EACf,mBAAmB,GACtB,MAAM,kBAAkB,CAAC;AAC1B,YAAY,EAAE,QAAQ,EAAE,MAAM,kBAAkB,CAAC;AAIjD,OAAO,EAAE,SAAS,EAAE,WAAW,EAAE,KAAK,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AAI9E,YAAY,EACR,sBAAsB,EACtB,gBAAgB,EAChB,kBAAkB,EAClB,eAAe,EACf,SAAS,GACZ,MAAM,wBAAwB,CAAC;AAEhC,YAAY,EAAE,aAAa,EAAE,UAAU,EAAE,SAAS,EAAE,kBAAkB,EAAE,MAAM,2BAA2B,CAAC;AAC1G,YAAY,EAAE,iBAAiB,EAAE,UAAU,EAAE,UAAU,EAAE,kBAAkB,EAAE,gBAAgB,EAAE,MAAM,yBAAyB,CAAC;AAC/H,YAAY,EAAE,iBAAiB,EAAE,WAAW,EAAE,MAAM,yBAAyB,CAAC;AAC9E,YAAY,EAAE,iBAAiB,EAAE,WAAW,EAAE,MAAO,yBAAyB,CAAC;AAC/E,YAAY,EAAE,oBAAoB,EAAE,cAAc,EAAE,MAAM,4BAA4B,CAAC;AAIvF,OAAO,EACH,UAAU,EACV,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,qBAAqB,EACrB,eAAe,EACf,cAAc,EACd,UAAU,EACV,YAAY,EACZ,IAAI,EACJ,UAAU,GACb,MAAM,sBAAsB,CAAC;AAI9B,OAAO,EACH,YAAY,EACZ,cAAc,EACd,WAAW,EACX,WAAW,EACX,YAAY,EACZ,cAAc,EACd,sBAAsB,EACtB,wBAAwB,EACxB,cAAc,GACjB,MAAM,yBAAyB,CAAC;AAKjC,OAAO,EAAE,2BAA2B,EAAE,4BAA4B,EAAE,MAC3D,6BAA6B,CAAC;AACvC,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,qBAAqB,CAAC;AAC/B,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,gCAAgC,CAAC;AAC1C,OAAO,EAAE,kBAAkB,EAAE,cAAc,EAAE,MACpC,4BAA4B,CAAC;AACtC,OAAO,EAAE,gBAAgB,EAAE,yBAAyB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,MAC1F,0BAA0B,CAAC;AAGpC,OAAO,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MACrC,kBAAkB,CAAC;AAG5B,OAAO,EAAE,wBAAwB,EAAE,yBAAyB,EAAE,MACrD,0BAA0B,CAAC;AAGpC,OAAO,EAAE,sBAAsB,EAAE,uBAAuB,EAAE,YAAY,EAAE,MAC/D,wBAAwB,CAAC;AAIlC,OAAO,EACH,MAAM,EACN,UAAU,EACV,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,QAAQ,EACR,UAAU,EACV,YAAY,EACZ,aAAa,EACb,aAAa,EACb,mBAAmB,EACnB,WAAW,EACX,qBAAqB,EACrB,mBAAmB,EACnB,aAAa,GAChB,MAAM,mBAAmB,CAAC;AAC3B,YAAY,EACR,MAAM,EACN,aAAa,EACb,iBAAiB,EACjB,iBAAiB,EACjB,aAAa,EACb,WAAW,EACX,YAAY,EACZ,kBAAkB,GACrB,MAAM,mBAAmB,CAAC;AAE3B,OAAO,EAAE,kBAAkB,EAAE,MAAM,4BAA4B,CAAC;AAIhE,eAAO,MAAM,OAAO,UAAc,CAAC;AACnC,eAAO,MAAM,WAAW,wFAAwF,CAAC"}
package/dist/index.js CHANGED
@@ -15,6 +15,9 @@ export { SharedExpertMoE, LoadBalanceAccumulator, DEFAULT_MOE_CONFIG, DEFAULT_MO
15
15
  // ── EvermindLM (the generative model) + AdamW ──────────────────────────────────
16
16
  export { EvermindLM, EvermindLMTrainer, DEFAULT_LM_CONFIG, DEFAULT_LM_SEED } from './lm/index.js';
17
17
  export { AdamW } from './optim/adamw.js';
18
+ // PEFT / efficient-training toolkit (LoRA, QLoRA, mixed precision, checkpointing).
19
+ export { LoRAAdapter, EvermindLMLoRA, quantizeBase } from './training/lora.js';
20
+ export { DynamicLossScaler, roundFp16, fp16View } from './training/mixed_precision.js';
18
21
  // ── Training ──────────────────────────────────────────────────────────────────
19
22
  export { MambaTrainer } from './training/trainer.js';
20
23
  export { Tensor, backward, enableGrad, noGrad, clearTape, recordOperation, crossEntropyLoss, crossEntropyGrad, } from './training/autograd.js';
package/dist/index.js.map CHANGED
@@ -1 +1 @@
1
- {"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;GAEG;AAEH,iFAAiF;AAEjF,OAAO,EAAE,gBAAgB,EAAE,UAAU,EAAE,MAAM,wBAAwB,CAAC;AAEtE,oBAAoB;AACpB,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,cAAc,EAAE,MAAM,4BAA4B,CAAC;AAE5D,sCAAsC;AACtC,OAAO,EAAE,UAAU,EAAE,MAAM,yBAAyB,CAAC;AAErD,kFAAkF;AAClF,OAAO,EACH,eAAe,EACf,sBAAsB,EACtB,kBAAkB,EAClB,gBAAgB,EAChB,UAAU,EACV,oBAAoB,GACvB,MAAM,gBAAgB,CAAC;AAgBxB,kFAAkF;AAClF,OAAO,EAAE,UAAU,EAAE,iBAAiB,EAAE,iBAAiB,EAAE,eAAe,EAAE,MAAM,eAAe,CAAC;AAElG,OAAO,EAAE,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAGzC,iFAAiF;AAEjF,OAAO,EAAE,YAAY,EAAE,MAAM,uBAAuB,CAAC;AACrD,OAAO,EACH,MAAM,EACN,QAAQ,EACR,UAAU,EACV,MAAM,EACN,SAAS,EACT,eAAe,EACf,gBAAgB,EAChB,gBAAgB,GACnB,MAAM,wBAAwB,CAAC;AAEhC,iFAAiF;AAEjF,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAGlD,mFAAmF;AAEnF,OAAO,EACH,eAAe,EACf,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,kBAAkB,EAClB,aAAa,EACb,aAAa,GAChB,MAAM,kBAAkB,CAAC;AAS1B,mFAAmF;AAEnF,OAAO,EACH,cAAc,EACd,cAAc,EACd,iBAAiB,EACjB,oBAAoB,EACpB,UAAU,EACV,UAAU,EACV,UAAU,EACV,oBAAoB,EACpB,aAAa,EACb,iBAAiB,EACjB,YAAY,EACZ,kBAAkB,EAClB,MAAM,EACN,UAAU,GACb,MAAM,mBAAmB,CAAC;AAY3B,oFAAoF;AAEpF,OAAO,EACH,oBAAoB,EACpB,cAAc,EACd,qBAAqB,EACrB,oBAAoB,GACvB,MAAM,mBAAmB,CAAC;AAG3B,iFAAiF;AAEjF,OAAO,EACH,cAAc,EACd,mBAAmB,EACnB,aAAa,EACb,aAAa,EACb,cAAc,EACd,iBAAiB,EACjB,iBAAiB,EACjB,MAAM,IAAI,WAAW,EACrB,WAAW,EACX,UAAU,EACV,YAAY,EACZ,GAAG,GACN,MAAM,kBAAkB,CAAC;AAW1B,iFAAiF;AAEjF,OAAO,EAAE,KAAK,EAAE,gBAAgB,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MAAM,kBAAkB,CAAC;AAEhG,OAAO,EACH,eAAe,EACf,aAAa,EACb,iBAAiB,EACjB,mBAAmB,EACnB,eAAe,EACf,mBAAmB,GACtB,MAAM,kBAAkB,CAAC;AAG1B,iFAAiF;AAEjF,OAAO,EAAE,SAAS,EAAE,WAAW,EAAE,KAAK,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AAkB9E,iFAAiF;AAEjF,OAAO,EACH,UAAU,EACV,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,qBAAqB,EACrB,eAAe,EACf,cAAc,EACd,UAAU,EACV,YAAY,EACZ,IAAI,EACJ,UAAU,GACb,MAAM,sBAAsB,CAAC;AAE9B,iFAAiF;AAEjF,OAAO,EACH,YAAY,EACZ,cAAc,EACd,WAAW,EACX,WAAW,EACX,YAAY,EACZ,cAAc,EACd,sBAAsB,EACtB,wBAAwB,EACxB,cAAc,GACjB,MAAM,yBAAyB,CAAC;AAEjC,iFAAiF;AAEjF,8BAA8B;AAC9B,OAAO,EAAE,2BAA2B,EAAE,4BAA4B,EAAE,MAC3D,6BAA6B,CAAC;AACvC,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,qBAAqB,CAAC;AAC/B,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,gCAAgC,CAAC;AAC1C,OAAO,EAAE,kBAAkB,EAAE,cAAc,EAAE,MACpC,4BAA4B,CAAC;AACtC,OAAO,EAAE,gBAAgB,EAAE,yBAAyB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,MAC1F,0BAA0B,CAAC;AAEpC,sBAAsB;AACtB,OAAO,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MACrC,kBAAkB,CAAC;AAE5B,8BAA8B;AAC9B,OAAO,EAAE,wBAAwB,EAAE,yBAAyB,EAAE,MACrD,0BAA0B,CAAC;AAEpC,oBAAoB;AACpB,OAAO,EAAE,sBAAsB,EAAE,uBAAuB,EAAE,YAAY,EAAE,MAC/D,wBAAwB,CAAC;AAElC,iFAAiF;AAEjF,OAAO,EACH,MAAM,EACN,UAAU,EACV,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,QAAQ,EACR,UAAU,EACV,YAAY,EACZ,aAAa,EACb,aAAa,EACb,mBAAmB,EACnB,WAAW,EACX,qBAAqB,EACrB,mBAAmB,EACnB,aAAa,GAChB,MAAM,mBAAmB,CAAC;AAY3B,OAAO,EAAE,kBAAkB,EAAE,MAAM,4BAA4B,CAAC;AAEhE,iFAAiF;AAEjF,MAAM,CAAC,MAAM,OAAO,GAAO,OAAO,CAAC;AACnC,MAAM,CAAC,MAAM,WAAW,GAAG,qFAAqF,CAAC"}
1
+ {"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;GAEG;AAEH,iFAAiF;AAEjF,OAAO,EAAE,gBAAgB,EAAE,UAAU,EAAE,MAAM,wBAAwB,CAAC;AAEtE,oBAAoB;AACpB,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,WAAW,EAAE,MAAQ,yBAAyB,CAAC;AACxD,OAAO,EAAE,cAAc,EAAE,MAAM,4BAA4B,CAAC;AAE5D,sCAAsC;AACtC,OAAO,EAAE,UAAU,EAAE,MAAM,yBAAyB,CAAC;AAErD,kFAAkF;AAClF,OAAO,EACH,eAAe,EACf,sBAAsB,EACtB,kBAAkB,EAClB,gBAAgB,EAChB,UAAU,EACV,oBAAoB,GACvB,MAAM,gBAAgB,CAAC;AAgBxB,kFAAkF;AAClF,OAAO,EAAE,UAAU,EAAE,iBAAiB,EAAE,iBAAiB,EAAE,eAAe,EAAE,MAAM,eAAe,CAAC;AAElG,OAAO,EAAE,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAGzC,mFAAmF;AACnF,OAAO,EAAE,WAAW,EAAE,cAAc,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAE/E,OAAO,EAAE,iBAAiB,EAAE,SAAS,EAAE,QAAQ,EAAE,MAAM,+BAA+B,CAAC;AAGvF,iFAAiF;AAEjF,OAAO,EAAE,YAAY,EAAE,MAAM,uBAAuB,CAAC;AACrD,OAAO,EACH,MAAM,EACN,QAAQ,EACR,UAAU,EACV,MAAM,EACN,SAAS,EACT,eAAe,EACf,gBAAgB,EAChB,gBAAgB,GACnB,MAAM,wBAAwB,CAAC;AAEhC,iFAAiF;AAEjF,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAGlD,mFAAmF;AAEnF,OAAO,EACH,eAAe,EACf,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,kBAAkB,EAClB,aAAa,EACb,aAAa,GAChB,MAAM,kBAAkB,CAAC;AAS1B,mFAAmF;AAEnF,OAAO,EACH,cAAc,EACd,cAAc,EACd,iBAAiB,EACjB,oBAAoB,EACpB,UAAU,EACV,UAAU,EACV,UAAU,EACV,oBAAoB,EACpB,aAAa,EACb,iBAAiB,EACjB,YAAY,EACZ,kBAAkB,EAClB,MAAM,EACN,UAAU,GACb,MAAM,mBAAmB,CAAC;AAY3B,oFAAoF;AAEpF,OAAO,EACH,oBAAoB,EACpB,cAAc,EACd,qBAAqB,EACrB,oBAAoB,GACvB,MAAM,mBAAmB,CAAC;AAG3B,iFAAiF;AAEjF,OAAO,EACH,cAAc,EACd,mBAAmB,EACnB,aAAa,EACb,aAAa,EACb,cAAc,EACd,iBAAiB,EACjB,iBAAiB,EACjB,MAAM,IAAI,WAAW,EACrB,WAAW,EACX,UAAU,EACV,YAAY,EACZ,GAAG,GACN,MAAM,kBAAkB,CAAC;AAW1B,iFAAiF;AAEjF,OAAO,EAAE,KAAK,EAAE,gBAAgB,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MAAM,kBAAkB,CAAC;AAEhG,OAAO,EACH,eAAe,EACf,aAAa,EACb,iBAAiB,EACjB,mBAAmB,EACnB,eAAe,EACf,mBAAmB,GACtB,MAAM,kBAAkB,CAAC;AAG1B,iFAAiF;AAEjF,OAAO,EAAE,SAAS,EAAE,WAAW,EAAE,KAAK,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AAkB9E,iFAAiF;AAEjF,OAAO,EACH,UAAU,EACV,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,qBAAqB,EACrB,eAAe,EACf,cAAc,EACd,UAAU,EACV,YAAY,EACZ,IAAI,EACJ,UAAU,GACb,MAAM,sBAAsB,CAAC;AAE9B,iFAAiF;AAEjF,OAAO,EACH,YAAY,EACZ,cAAc,EACd,WAAW,EACX,WAAW,EACX,YAAY,EACZ,cAAc,EACd,sBAAsB,EACtB,wBAAwB,EACxB,cAAc,GACjB,MAAM,yBAAyB,CAAC;AAEjC,iFAAiF;AAEjF,8BAA8B;AAC9B,OAAO,EAAE,2BAA2B,EAAE,4BAA4B,EAAE,MAC3D,6BAA6B,CAAC;AACvC,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,qBAAqB,CAAC;AAC/B,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,gCAAgC,CAAC;AAC1C,OAAO,EAAE,kBAAkB,EAAE,cAAc,EAAE,MACpC,4BAA4B,CAAC;AACtC,OAAO,EAAE,gBAAgB,EAAE,yBAAyB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,MAC1F,0BAA0B,CAAC;AAEpC,sBAAsB;AACtB,OAAO,EAAE,gBAAgB,EAAE,iBAAiB,EAAE,MACrC,kBAAkB,CAAC;AAE5B,8BAA8B;AAC9B,OAAO,EAAE,wBAAwB,EAAE,yBAAyB,EAAE,MACrD,0BAA0B,CAAC;AAEpC,oBAAoB;AACpB,OAAO,EAAE,sBAAsB,EAAE,uBAAuB,EAAE,YAAY,EAAE,MAC/D,wBAAwB,CAAC;AAElC,iFAAiF;AAEjF,OAAO,EACH,MAAM,EACN,UAAU,EACV,gBAAgB,EAChB,gBAAgB,EAChB,aAAa,EACb,aAAa,EACb,QAAQ,EACR,UAAU,EACV,YAAY,EACZ,aAAa,EACb,aAAa,EACb,mBAAmB,EACnB,WAAW,EACX,qBAAqB,EACrB,mBAAmB,EACnB,aAAa,GAChB,MAAM,mBAAmB,CAAC;AAY3B,OAAO,EAAE,kBAAkB,EAAE,MAAM,4BAA4B,CAAC;AAEhE,iFAAiF;AAEjF,MAAM,CAAC,MAAM,OAAO,GAAO,OAAO,CAAC;AACnC,MAAM,CAAC,MAAM,WAAW,GAAG,qFAAqF,CAAC"}
@@ -21,6 +21,7 @@
21
21
  * acceleration with the same shapes.
22
22
  */
23
23
  import { type AdamWOptions } from "../optim/adamw.js";
24
+ import { DynamicLossScaler, type LossScalerOptions } from "../training/mixed_precision.js";
24
25
  export interface EvermindLMConfig {
25
26
  /** Vocabulary size (the only required field; everything else has a default). */
26
27
  vocabSize: number;
@@ -97,17 +98,53 @@ export declare class EvermindLM {
97
98
  /** Per-layer channel mixer. */
98
99
  private readonly moe;
99
100
  constructor(config: EvermindLMConfig);
101
+ /** Embed a token sequence into per-position channel vectors. */
102
+ private _embed;
103
+ /**
104
+ * One (conv + MoE) block: pre-norm → depthwise causal conv → residual, then
105
+ * pre-norm → MoE channel mixer → residual. Returns the block output and the
106
+ * activation cache its backward needs. Isolating this is what lets
107
+ * {@link lossAndBackwardCheckpointed} recompute a layer's activations on demand
108
+ * instead of retaining every layer's cache at once.
109
+ */
110
+ private _forwardLayer;
111
+ /** Tied output head: logits_t[v] = x_t · emb[v]. */
112
+ private _head;
100
113
  /** Run the model over a token sequence; returns per-position logits + a cache. */
101
114
  forward(tokens: number[]): {
102
115
  logits: Float32Array[];
103
116
  cache: ForwardCache;
104
117
  };
118
+ /**
119
+ * Head + tied-embedding gradient. Accumulates dL/d(head→emb) into gEmb and
120
+ * returns the mean next-token loss plus dL/d(finalX). Shared by the full and
121
+ * checkpointed backward paths so the head maths lives in one place.
122
+ */
123
+ private _headBackward;
124
+ /**
125
+ * Backward through one (conv + MoE) block given dL/d(block output) and the
126
+ * block's activation cache. Accumulates conv/norm/MoE gradients and returns
127
+ * dL/d(block input). The inverse of {@link _forwardLayer}.
128
+ */
129
+ private _backwardLayer;
130
+ /** Embedding lookup: dL/d(layer-0 input) flows into the row for each token. */
131
+ private _embedBackward;
105
132
  /**
106
133
  * Next-token cross-entropy over the sequence (predict tokens[t+1] from
107
134
  * position t), accumulating exact gradients. Returns the mean loss. Call
108
135
  * {@link zeroGrad} before and an optimiser step after.
109
136
  */
110
137
  lossAndBackward(tokens: number[]): number;
138
+ /**
139
+ * Activation-checkpointed backward — numerically identical gradients to
140
+ * {@link lossAndBackward}, but retains only the per-LAYER inputs during the
141
+ * forward instead of every layer's full activation cache. Each layer's
142
+ * activations are RECOMPUTED (a cheap extra forward) when its backward runs,
143
+ * so peak activation memory is one layer's cache, not all of them. This is the
144
+ * memory-for-compute trade the cookbook pairs with FSDP to fit longer
145
+ * sequences / bigger models on a constrained device.
146
+ */
147
+ lossAndBackwardCheckpointed(tokens: number[]): number;
111
148
  /**
112
149
  * Text-level generation: encode the prompt, generate, decode. `codec` is any
113
150
  * tokenizer exposing encode/decode (the engine's `BPETokenizer` satisfies it),
@@ -149,14 +186,44 @@ export declare class EvermindLM {
149
186
  /** Reconstruct weights from a base EVL0 checkpoint + a delta (EVM-6). */
150
187
  loadDelta(baseCheckpoint: ArrayBuffer, delta: ArrayBuffer): void;
151
188
  }
189
+ export interface EvermindLMTrainOptions extends AdamWOptions {
190
+ epochs?: number;
191
+ /**
192
+ * Gradient accumulation: average gradients over this many sequences (micro-
193
+ * batches) before each optimiser step, for a larger effective batch on a
194
+ * memory-constrained device. Default 1.
195
+ */
196
+ accumSteps?: number;
197
+ /**
198
+ * Use activation checkpointing (recompute layer activations in backward) to
199
+ * cap peak activation memory. Identical gradients, a little extra compute.
200
+ * Default false.
201
+ */
202
+ checkpoint?: boolean;
203
+ /**
204
+ * Mixed-precision training: fp16-rounded gradients with dynamic loss scaling
205
+ * over fp32 master weights (the model params). Overflowing steps are skipped
206
+ * and the scale backs off. Default false.
207
+ */
208
+ mixedPrecision?: boolean | LossScalerOptions;
209
+ }
152
210
  /** Minimal sequence trainer: AdamW over next-token cross-entropy. */
153
211
  export declare class EvermindLMTrainer {
154
212
  private readonly model;
155
213
  private readonly opts;
156
214
  private readonly adam;
157
- constructor(model: EvermindLM, opts?: AdamWOptions & {
158
- epochs?: number;
159
- });
215
+ private readonly scaler;
216
+ private readonly grads;
217
+ constructor(model: EvermindLM, opts?: EvermindLMTrainOptions);
218
+ /** The dynamic loss scaler (mixed-precision mode only) — exposes scale/overflow stats. */
219
+ get lossScaler(): DynamicLossScaler | null;
220
+ private _backward;
221
+ /** Divide accumulated gradients by `n` (accumulation averaging), in place. */
222
+ private _scaleGrads;
223
+ /** Round accumulated gradients to fp16 precision (mixed-precision simulation), in place. */
224
+ private _fp16Grads;
225
+ /** Multiply accumulated gradients by `s`, in place. */
226
+ private _mulGrads;
160
227
  /** Train on a set of token sequences; returns per-epoch mean loss. */
161
228
  fit(sequences: number[][]): number[];
162
229
  }
@@ -1 +1 @@
1
- {"version":3,"file":"evermind_lm.d.ts","sourceRoot":"","sources":["../../src/lm/evermind_lm.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;GAqBG;AAIH,OAAO,EAAS,KAAK,YAAY,EAAE,MAAM,mBAAmB,CAAC;AAM7D,MAAM,WAAW,gBAAgB;IAC/B,gFAAgF;IAChF,SAAS,EAAE,MAAM,CAAC;IAClB,6CAA6C;IAC7C,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,gDAAgD;IAChD,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,2CAA2C;IAC3C,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,6DAA6D;IAC7D,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,+CAA+C;IAC/C,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,8CAA8C;IAC9C,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,+BAA+B;IAC/B,IAAI,CAAC,EAAE,MAAM,CAAC;CACf;AAED,eAAO,MAAM,iBAAiB,EAAE,QAAQ,CAAC,IAAI,CAAC,gBAAgB,EAAE,MAAM,GAAG,WAAW,CAAC,CAOpF,CAAC;AAEF,eAAO,MAAM,eAAe,aAAa,CAAC;AAG1C,UAAU,YAAY;IACpB,CAAC,EAAE,YAAY,CAAC;IAChB,KAAK,EAAE;QAAE,OAAO,EAAE,MAAM,EAAE,CAAC;QAAC,KAAK,EAAE,MAAM,EAAE,CAAC;QAAC,KAAK,EAAE,YAAY,CAAA;KAAE,CAAC;IACnE,SAAS,EAAE,YAAY,CAAC;IACxB,OAAO,EAAE,YAAY,CAAC;IACtB,SAAS,EAAE,YAAY,EAAE,CAAC;IAC1B,SAAS,EAAE,YAAY,EAAE,CAAC;IAC1B,OAAO,EAAE,YAAY,EAAE,CAAC;CACzB;AAED,UAAU,UAAU;IAClB,OAAO,EAAE,YAAY,EAAE,CAAC;IACxB,UAAU,EAAE,YAAY,EAAE,CAAC;IAC3B,OAAO,EAAE,MAAM,EAAE,CAAC;IAClB,SAAS,EAAE,YAAY,EAAE,CAAC;IAC1B,MAAM,EAAE,MAAM,EAAE,CAAC;IACjB,QAAQ,EAAE,YAAY,EAAE,CAAC;CAC1B;AAED,UAAU,YAAY;IACpB,MAAM,EAAE,MAAM,EAAE,CAAC;IACjB,MAAM,EAAE,UAAU,EAAE,CAAC;IACrB,MAAM,EAAE,YAAY,EAAE,CAAC;CACxB;AAED,yFAAyF;AACzF,MAAM,WAAW,SAAS;IACxB,MAAM,CAAC,IAAI,EAAE,MAAM,GAAG,MAAM,EAAE,CAAC;IAC/B,MAAM,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,MAAM,CAAC;CAC/B;AAED,MAAM,WAAW,iBAAiB;IAChC,YAAY,EAAE,MAAM,CAAC;IACrB,oEAAoE;IACpE,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,mEAAmE;IACnE,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,sDAAsD;IACtD,SAAS,CAAC,EAAE,MAAM,CAAC;CACpB;AAED,qBAAa,UAAU;IACrB,QAAQ,CAAC,MAAM,EAAE,QAAQ,CAAC,IAAI,CAAC,gBAAgB,EAAE,MAAM,CAAC,CAAC,CAAC;IAE1D,0EAA0E;IAC1E,GAAG,EAAE,YAAY,CAAC;IAClB,OAAO,CAAC,IAAI,CAAe;IAC3B,oEAAoE;IACpE,OAAO,CAAC,QAAQ,CAAC,IAAI,CAAiB;IACtC,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAiB;IACvC,gEAAgE;IAChE,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAiB;IACvC,OAAO,CAAC,QAAQ,CAAC,MAAM,CAAiB;IACxC,OAAO,CAAC,QAAQ,CAAC,IAAI,CAAiB;IACtC,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAiB;IACvC,+BAA+B;IAC/B,OAAO,CAAC,QAAQ,CAAC,GAAG,CAAoB;gBAE5B,MAAM,EAAE,gBAAgB;IA+DpC,kFAAkF;IAClF,OAAO,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG;QAAE,MAAM,EAAE,YAAY,EAAE,CAAC;QAAC,KAAK,EAAE,YAAY,CAAA;KAAE;IA4E1E;;;;OAIG;IACH,eAAe,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG,MAAM;IAqFzC;;;;;OAKG;IACH,YAAY,CAAC,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,SAAS,EAAE,IAAI,EAAE,iBAAiB,GAAG,MAAM;IAI/E,qFAAqF;IACrF,QAAQ,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,IAAI,EAAE,iBAAiB,GAAG,MAAM,EAAE;IAkB7D,8EAA8E;IAC9E,UAAU,IAAI;QAAE,IAAI,EAAE,YAAY,CAAA;KAAE,EAAE;IAStC,+DAA+D;IAC/D,SAAS,IAAI;QAAE,IAAI,EAAE,YAAY,CAAA;KAAE,EAAE;IASrC,QAAQ,IAAI,IAAI;IAUhB,uFAAuF;IACvF,aAAa,CAAC,IAAI,GAAE;QAAE,IAAI,CAAC,EAAE,OAAO,CAAA;KAAO,GAAG,WAAW;IAiCzD,sFAAsF;IACtF,WAAW,CAAC,MAAM,EAAE,WAAW,GAAG,IAAI;IAItC,oFAAoF;IACpF,OAAO,CAAC,SAAS;IA0BjB,qFAAqF;IACrF,OAAO,CAAC,KAAK;IASb,uEAAuE;IACvE,OAAO,CAAC,QAAQ;IAQhB;;;;OAIG;IACH,WAAW,CAAC,cAAc,EAAE,WAAW,EAAE,IAAI,GAAE;QAAE,GAAG,CAAC,EAAE,MAAM,CAAA;KAAO,GAAG,WAAW;IASlF,yEAAyE;IACzE,SAAS,CAAC,cAAc,EAAE,WAAW,EAAE,KAAK,EAAE,WAAW,GAAG,IAAI;CAIjE;AAED,qEAAqE;AACrE,qBAAa,iBAAiB;IAG1B,OAAO,CAAC,QAAQ,CAAC,KAAK;IACtB,OAAO,CAAC,QAAQ,CAAC,IAAI;IAHvB,OAAO,CAAC,QAAQ,CAAC,IAAI,CAAQ;gBAEV,KAAK,EAAE,UAAU,EACjB,IAAI,GAAE,YAAY,GAAG;QAAE,MAAM,CAAC,EAAE,MAAM,CAAA;KAAO;IAIhE,sEAAsE;IACtE,GAAG,CAAC,SAAS,EAAE,MAAM,EAAE,EAAE,GAAG,MAAM,EAAE;CAiBrC"}
1
+ {"version":3,"file":"evermind_lm.d.ts","sourceRoot":"","sources":["../../src/lm/evermind_lm.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;GAqBG;AAIH,OAAO,EAAS,KAAK,YAAY,EAAE,MAAM,mBAAmB,CAAC;AAC7D,OAAO,EAAE,iBAAiB,EAAa,KAAK,iBAAiB,EAAE,MAAM,gCAAgC,CAAC;AAMtG,MAAM,WAAW,gBAAgB;IAC/B,gFAAgF;IAChF,SAAS,EAAE,MAAM,CAAC;IAClB,6CAA6C;IAC7C,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,gDAAgD;IAChD,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,2CAA2C;IAC3C,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,6DAA6D;IAC7D,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,+CAA+C;IAC/C,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,8CAA8C;IAC9C,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,+BAA+B;IAC/B,IAAI,CAAC,EAAE,MAAM,CAAC;CACf;AAED,eAAO,MAAM,iBAAiB,EAAE,QAAQ,CAAC,IAAI,CAAC,gBAAgB,EAAE,MAAM,GAAG,WAAW,CAAC,CAOpF,CAAC;AAEF,eAAO,MAAM,eAAe,aAAa,CAAC;AAG1C,UAAU,YAAY;IACpB,CAAC,EAAE,YAAY,CAAC;IAChB,KAAK,EAAE;QAAE,OAAO,EAAE,MAAM,EAAE,CAAC;QAAC,KAAK,EAAE,MAAM,EAAE,CAAC;QAAC,KAAK,EAAE,YAAY,CAAA;KAAE,CAAC;IACnE,SAAS,EAAE,YAAY,CAAC;IACxB,OAAO,EAAE,YAAY,CAAC;IACtB,SAAS,EAAE,YAAY,EAAE,CAAC;IAC1B,SAAS,EAAE,YAAY,EAAE,CAAC;IAC1B,OAAO,EAAE,YAAY,EAAE,CAAC;CACzB;AAED,UAAU,UAAU;IAClB,OAAO,EAAE,YAAY,EAAE,CAAC;IACxB,UAAU,EAAE,YAAY,EAAE,CAAC;IAC3B,OAAO,EAAE,MAAM,EAAE,CAAC;IAClB,SAAS,EAAE,YAAY,EAAE,CAAC;IAC1B,MAAM,EAAE,MAAM,EAAE,CAAC;IACjB,QAAQ,EAAE,YAAY,EAAE,CAAC;CAC1B;AAED,UAAU,YAAY;IACpB,MAAM,EAAE,MAAM,EAAE,CAAC;IACjB,MAAM,EAAE,UAAU,EAAE,CAAC;IACrB,MAAM,EAAE,YAAY,EAAE,CAAC;CACxB;AAED,yFAAyF;AACzF,MAAM,WAAW,SAAS;IACxB,MAAM,CAAC,IAAI,EAAE,MAAM,GAAG,MAAM,EAAE,CAAC;IAC/B,MAAM,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,MAAM,CAAC;CAC/B;AAED,MAAM,WAAW,iBAAiB;IAChC,YAAY,EAAE,MAAM,CAAC;IACrB,oEAAoE;IACpE,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,mEAAmE;IACnE,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,sDAAsD;IACtD,SAAS,CAAC,EAAE,MAAM,CAAC;CACpB;AAED,qBAAa,UAAU;IACrB,QAAQ,CAAC,MAAM,EAAE,QAAQ,CAAC,IAAI,CAAC,gBAAgB,EAAE,MAAM,CAAC,CAAC,CAAC;IAE1D,0EAA0E;IAC1E,GAAG,EAAE,YAAY,CAAC;IAClB,OAAO,CAAC,IAAI,CAAe;IAC3B,oEAAoE;IACpE,OAAO,CAAC,QAAQ,CAAC,IAAI,CAAiB;IACtC,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAiB;IACvC,gEAAgE;IAChE,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAiB;IACvC,OAAO,CAAC,QAAQ,CAAC,MAAM,CAAiB;IACxC,OAAO,CAAC,QAAQ,CAAC,IAAI,CAAiB;IACtC,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAiB;IACvC,+BAA+B;IAC/B,OAAO,CAAC,QAAQ,CAAC,GAAG,CAAoB;gBAE5B,MAAM,EAAE,gBAAgB;IA+DpC,gEAAgE;IAChE,OAAO,CAAC,MAAM;IAUd;;;;;;OAMG;IACH,OAAO,CAAC,aAAa;IA2CrB,oDAAoD;IACpD,OAAO,CAAC,KAAK;IAcb,kFAAkF;IAClF,OAAO,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG;QAAE,MAAM,EAAE,YAAY,EAAE,CAAC;QAAC,KAAK,EAAE,YAAY,CAAA;KAAE;IAa1E;;;;OAIG;IACH,OAAO,CAAC,aAAa;IAyBrB;;;;OAIG;IACH,OAAO,CAAC,cAAc;IA8CtB,+EAA+E;IAC/E,OAAO,CAAC,cAAc;IAQtB;;;;OAIG;IACH,eAAe,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG,MAAM;IAUzC;;;;;;;;OAQG;IACH,2BAA2B,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG,MAAM;IAqBrD;;;;;OAKG;IACH,YAAY,CAAC,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,SAAS,EAAE,IAAI,EAAE,iBAAiB,GAAG,MAAM;IAI/E,qFAAqF;IACrF,QAAQ,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,IAAI,EAAE,iBAAiB,GAAG,MAAM,EAAE;IAkB7D,8EAA8E;IAC9E,UAAU,IAAI;QAAE,IAAI,EAAE,YAAY,CAAA;KAAE,EAAE;IAStC,+DAA+D;IAC/D,SAAS,IAAI;QAAE,IAAI,EAAE,YAAY,CAAA;KAAE,EAAE;IASrC,QAAQ,IAAI,IAAI;IAUhB,uFAAuF;IACvF,aAAa,CAAC,IAAI,GAAE;QAAE,IAAI,CAAC,EAAE,OAAO,CAAA;KAAO,GAAG,WAAW;IAiCzD,sFAAsF;IACtF,WAAW,CAAC,MAAM,EAAE,WAAW,GAAG,IAAI;IAItC,oFAAoF;IACpF,OAAO,CAAC,SAAS;IA0BjB,qFAAqF;IACrF,OAAO,CAAC,KAAK;IASb,uEAAuE;IACvE,OAAO,CAAC,QAAQ;IAQhB;;;;OAIG;IACH,WAAW,CAAC,cAAc,EAAE,WAAW,EAAE,IAAI,GAAE;QAAE,GAAG,CAAC,EAAE,MAAM,CAAA;KAAO,GAAG,WAAW;IASlF,yEAAyE;IACzE,SAAS,CAAC,cAAc,EAAE,WAAW,EAAE,KAAK,EAAE,WAAW,GAAG,IAAI;CAIjE;AAED,MAAM,WAAW,sBAAuB,SAAQ,YAAY;IAC1D,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB;;;;OAIG;IACH,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB;;;;OAIG;IACH,UAAU,CAAC,EAAE,OAAO,CAAC;IACrB;;;;OAIG;IACH,cAAc,CAAC,EAAE,OAAO,GAAG,iBAAiB,CAAC;CAC9C;AAED,qEAAqE;AACrE,qBAAa,iBAAiB;IAK1B,OAAO,CAAC,QAAQ,CAAC,KAAK;IACtB,OAAO,CAAC,QAAQ,CAAC,IAAI;IALvB,OAAO,CAAC,QAAQ,CAAC,IAAI,CAAQ;IAC7B,OAAO,CAAC,QAAQ,CAAC,MAAM,CAA2B;IAClD,OAAO,CAAC,QAAQ,CAAC,KAAK,CAA2B;gBAE9B,KAAK,EAAE,UAAU,EACjB,IAAI,GAAE,sBAA2B;IASpD,0FAA0F;IAC1F,IAAI,UAAU,IAAI,iBAAiB,GAAG,IAAI,CAEzC;IAED,OAAO,CAAC,SAAS;IAIjB,8EAA8E;IAC9E,OAAO,CAAC,WAAW;IAMnB,4FAA4F;IAC5F,OAAO,CAAC,UAAU;IAIlB,uDAAuD;IACvD,OAAO,CAAC,SAAS;IAKjB,sEAAsE;IACtE,GAAG,CAAC,SAAS,EAAE,MAAM,EAAE,EAAE,GAAG,MAAM,EAAE;CA8CrC"}
@@ -23,6 +23,7 @@
23
23
  import { SharedExpertMoE } from "../moe/moe_model.js";
24
24
  import { crossEntropyLoss, crossEntropyGrad } from "../training/autograd.js";
25
25
  import { AdamW } from "../optim/adamw.js";
26
+ import { DynamicLossScaler, roundFp16 } from "../training/mixed_precision.js";
26
27
  import { SeededRng } from "../utils/rng.js";
27
28
  import { quantizeFp16, dequantizeFp16 } from "../utils/quantization.js";
28
29
  import { appendCrcTrailer, verifyCrcTrailer } from "../utils/crc32.js";
@@ -112,65 +113,70 @@ export class EvermindLM {
112
113
  }
113
114
  }
114
115
  // ── Forward ────────────────────────────────────────────────────────────────
115
- /** Run the model over a token sequence; returns per-position logits + a cache. */
116
- forward(tokens) {
117
- const { dModel, convKernel, numLayers, vocabSize } = this.config;
118
- const T = tokens.length;
119
- // Embed.
120
- let x = tokens.map((tok) => {
116
+ /** Embed a token sequence into per-position channel vectors. */
117
+ _embed(tokens) {
118
+ const { dModel } = this.config;
119
+ return tokens.map((tok) => {
121
120
  const row = new Float32Array(dModel);
122
121
  const off = tok * dModel;
123
122
  for (let c = 0; c < dModel; c++)
124
123
  row[c] = this.emb[off + c];
125
124
  return row;
126
125
  });
127
- const layers = [];
128
- for (let l = 0; l < numLayers; l++) {
129
- const layerIn = x;
130
- const ker = this.conv[l];
131
- const nConv = this.nConv[l];
132
- const nMoe = this.nMoe[l];
133
- // Pre-norm → depthwise causal conv → residual.
134
- const normedConv = [];
135
- const rmsConv = [];
136
- for (let t = 0; t < T; t++) {
137
- const { y, r } = rmsNorm(layerIn[t], nConv);
138
- normedConv.push(y);
139
- rmsConv.push(r);
140
- }
141
- const afterConv = [];
142
- for (let t = 0; t < T; t++) {
143
- const out = Float32Array.from(layerIn[t]); // residual base
144
- for (let c = 0; c < dModel; c++) {
145
- let acc = 0;
146
- for (let j = 0; j < convKernel; j++) {
147
- const ti = t - j;
148
- if (ti >= 0)
149
- acc += ker[c * convKernel + j] * normedConv[ti][c];
150
- }
151
- out[c] = out[c] + acc;
126
+ }
127
+ /**
128
+ * One (conv + MoE) block: pre-norm → depthwise causal conv → residual, then
129
+ * pre-norm → MoE channel mixer → residual. Returns the block output and the
130
+ * activation cache its backward needs. Isolating this is what lets
131
+ * {@link lossAndBackwardCheckpointed} recompute a layer's activations on demand
132
+ * instead of retaining every layer's cache at once.
133
+ */
134
+ _forwardLayer(l, layerIn) {
135
+ const { dModel, convKernel } = this.config;
136
+ const T = layerIn.length;
137
+ const ker = this.conv[l];
138
+ const nConv = this.nConv[l];
139
+ const nMoe = this.nMoe[l];
140
+ const normedConv = [];
141
+ const rmsConv = [];
142
+ for (let t = 0; t < T; t++) {
143
+ const { y, r } = rmsNorm(layerIn[t], nConv);
144
+ normedConv.push(y);
145
+ rmsConv.push(r);
146
+ }
147
+ const afterConv = [];
148
+ for (let t = 0; t < T; t++) {
149
+ const out = Float32Array.from(layerIn[t]); // residual base
150
+ for (let c = 0; c < dModel; c++) {
151
+ let acc = 0;
152
+ for (let j = 0; j < convKernel; j++) {
153
+ const ti = t - j;
154
+ if (ti >= 0)
155
+ acc += ker[c * convKernel + j] * normedConv[ti][c];
152
156
  }
153
- afterConv.push(out);
154
- }
155
- // Pre-norm → MoE channel mixer → residual.
156
- const rmsMoe = [];
157
- const moeCache = [];
158
- const afterMoe = [];
159
- for (let t = 0; t < T; t++) {
160
- const { y, r } = rmsNorm(afterConv[t], nMoe);
161
- rmsMoe.push(r);
162
- const out = Float32Array.from(afterConv[t]); // residual base
163
- const mr = this.moe[l].forward(y);
164
- for (let c = 0; c < dModel; c++)
165
- out[c] = out[c] + mr.output[c];
166
- afterMoe.push(out);
167
- moeCache.push(mr.cache);
157
+ out[c] = out[c] + acc;
168
158
  }
169
- layers.push({ layerIn, normedConv, rmsConv, afterConv, rmsMoe, moeCache });
170
- x = afterMoe;
159
+ afterConv.push(out);
171
160
  }
172
- // Tied head: logits_t[v] = x_t · emb[v].
173
- const logits = x.map((xt) => {
161
+ const rmsMoe = [];
162
+ const moeCache = [];
163
+ const afterMoe = [];
164
+ for (let t = 0; t < T; t++) {
165
+ const { y, r } = rmsNorm(afterConv[t], nMoe);
166
+ rmsMoe.push(r);
167
+ const out = Float32Array.from(afterConv[t]); // residual base
168
+ const mr = this.moe[l].forward(y);
169
+ for (let c = 0; c < dModel; c++)
170
+ out[c] = out[c] + mr.output[c];
171
+ afterMoe.push(out);
172
+ moeCache.push(mr.cache);
173
+ }
174
+ return { afterMoe, cache: { layerIn, normedConv, rmsConv, afterConv, rmsMoe, moeCache } };
175
+ }
176
+ /** Tied output head: logits_t[v] = x_t · emb[v]. */
177
+ _head(x) {
178
+ const { dModel, vocabSize } = this.config;
179
+ return x.map((xt) => {
174
180
  const lg = new Float32Array(vocabSize);
175
181
  for (let v = 0; v < vocabSize; v++) {
176
182
  let acc = 0;
@@ -181,30 +187,36 @@ export class EvermindLM {
181
187
  }
182
188
  return lg;
183
189
  });
184
- return { logits, cache: { tokens, layers, finalX: x } };
190
+ }
191
+ /** Run the model over a token sequence; returns per-position logits + a cache. */
192
+ forward(tokens) {
193
+ let x = this._embed(tokens);
194
+ const layers = [];
195
+ for (let l = 0; l < this.config.numLayers; l++) {
196
+ const { afterMoe, cache } = this._forwardLayer(l, x);
197
+ layers.push(cache);
198
+ x = afterMoe;
199
+ }
200
+ return { logits: this._head(x), cache: { tokens, layers, finalX: x } };
185
201
  }
186
202
  // ── Loss + backward ──────────────────────────────────────────────────────────
187
203
  /**
188
- * Next-token cross-entropy over the sequence (predict tokens[t+1] from
189
- * position t), accumulating exact gradients. Returns the mean loss. Call
190
- * {@link zeroGrad} before and an optimiser step after.
204
+ * Head + tied-embedding gradient. Accumulates dL/d(head→emb) into gEmb and
205
+ * returns the mean next-token loss plus dL/d(finalX). Shared by the full and
206
+ * checkpointed backward paths so the head maths lives in one place.
191
207
  */
192
- lossAndBackward(tokens) {
193
- const { dModel, convKernel, numLayers, vocabSize } = this.config;
208
+ _headBackward(tokens, logits, finalX) {
209
+ const { dModel, vocabSize } = this.config;
194
210
  const T = tokens.length;
195
- if (T < 2)
196
- return 0;
197
- const { logits, cache } = this.forward(tokens);
198
211
  const predPositions = T - 1; // positions 0..T-2 predict the next token
199
212
  const inv = 1 / predPositions;
200
- // dL/d(finalX_t) and head gradient into the tied embedding.
201
213
  const dX = Array.from({ length: T }, () => new Float32Array(dModel));
202
214
  let loss = 0;
203
215
  for (let t = 0; t < predPositions; t++) {
204
216
  const target = tokens[t + 1];
205
217
  loss += crossEntropyLoss(logits[t], target) * inv;
206
218
  const dLogit = crossEntropyGrad(logits[t], target); // probs - onehot
207
- const xt = cache.finalX[t];
219
+ const xt = finalX[t];
208
220
  for (let v = 0; v < vocabSize; v++) {
209
221
  const g = dLogit[v] * inv;
210
222
  if (g === 0)
@@ -216,60 +228,111 @@ export class EvermindLM {
216
228
  }
217
229
  }
218
230
  }
219
- // Backprop through layers in reverse.
220
- for (let l = numLayers - 1; l >= 0; l--) {
221
- const lc = cache.layers[l];
222
- const ker = this.conv[l];
223
- const gker = this.gConv[l];
224
- const nConv = this.nConv[l];
225
- const gNConv = this.gNConv[l];
226
- const nMoe = this.nMoe[l];
227
- const gNMoe = this.gNMoe[l];
228
- // MoE sub-block: afterMoe = afterConv + MoE(RMSNorm(afterConv, nMoe)).
229
- const dAfterConv = [];
230
- for (let t = 0; t < T; t++) {
231
- const dMoeNormed = this.moe[l].backward(dX[t], lc.moeCache[t]);
232
- const { dx, dgain } = rmsNormBackward(dMoeNormed, lc.afterConv[t], lc.rmsMoe[t], nMoe);
233
- for (let c = 0; c < dModel; c++)
234
- gNMoe[c] = gNMoe[c] + dgain[c];
235
- const d = Float32Array.from(dX[t]); // residual passthrough
236
- for (let c = 0; c < dModel; c++)
237
- d[c] = d[c] + dx[c];
238
- dAfterConv.push(d);
239
- }
240
- // Conv sub-block: afterConv = layerIn + conv(RMSNorm(layerIn, nConv)).
241
- const dNormedConv = Array.from({ length: T }, () => new Float32Array(dModel));
242
- const dLayerIn = dAfterConv.map((v) => Float32Array.from(v)); // residual passthrough
243
- for (let t = 0; t < T; t++) {
244
- for (let c = 0; c < dModel; c++) {
245
- const dmix = dAfterConv[t][c];
246
- if (dmix === 0)
231
+ return { loss, dX };
232
+ }
233
+ /**
234
+ * Backward through one (conv + MoE) block given dL/d(block output) and the
235
+ * block's activation cache. Accumulates conv/norm/MoE gradients and returns
236
+ * dL/d(block input). The inverse of {@link _forwardLayer}.
237
+ */
238
+ _backwardLayer(l, dOut, lc) {
239
+ const { dModel, convKernel } = this.config;
240
+ const T = dOut.length;
241
+ const ker = this.conv[l];
242
+ const gker = this.gConv[l];
243
+ const nConv = this.nConv[l];
244
+ const gNConv = this.gNConv[l];
245
+ const nMoe = this.nMoe[l];
246
+ const gNMoe = this.gNMoe[l];
247
+ // MoE sub-block: afterMoe = afterConv + MoE(RMSNorm(afterConv, nMoe)).
248
+ const dAfterConv = [];
249
+ for (let t = 0; t < T; t++) {
250
+ const dMoeNormed = this.moe[l].backward(dOut[t], lc.moeCache[t]);
251
+ const { dx, dgain } = rmsNormBackward(dMoeNormed, lc.afterConv[t], lc.rmsMoe[t], nMoe);
252
+ for (let c = 0; c < dModel; c++)
253
+ gNMoe[c] = gNMoe[c] + dgain[c];
254
+ const d = Float32Array.from(dOut[t]); // residual passthrough
255
+ for (let c = 0; c < dModel; c++)
256
+ d[c] = d[c] + dx[c];
257
+ dAfterConv.push(d);
258
+ }
259
+ // Conv sub-block: afterConv = layerIn + conv(RMSNorm(layerIn, nConv)).
260
+ const dNormedConv = Array.from({ length: T }, () => new Float32Array(dModel));
261
+ const dLayerIn = dAfterConv.map((v) => Float32Array.from(v)); // residual passthrough
262
+ for (let t = 0; t < T; t++) {
263
+ for (let c = 0; c < dModel; c++) {
264
+ const dmix = dAfterConv[t][c];
265
+ if (dmix === 0)
266
+ continue;
267
+ for (let j = 0; j < convKernel; j++) {
268
+ const ti = t - j;
269
+ if (ti < 0)
247
270
  continue;
248
- for (let j = 0; j < convKernel; j++) {
249
- const ti = t - j;
250
- if (ti < 0)
251
- continue;
252
- gker[c * convKernel + j] = gker[c * convKernel + j] + dmix * lc.normedConv[ti][c];
253
- dNormedConv[ti][c] = dNormedConv[ti][c] + dmix * ker[c * convKernel + j];
254
- }
255
- }
256
- }
257
- for (let t = 0; t < T; t++) {
258
- const { dx, dgain } = rmsNormBackward(dNormedConv[t], lc.layerIn[t], lc.rmsConv[t], nConv);
259
- for (let c = 0; c < dModel; c++) {
260
- gNConv[c] = gNConv[c] + dgain[c];
261
- dLayerIn[t][c] = dLayerIn[t][c] + dx[c];
271
+ gker[c * convKernel + j] = gker[c * convKernel + j] + dmix * lc.normedConv[ti][c];
272
+ dNormedConv[ti][c] = dNormedConv[ti][c] + dmix * ker[c * convKernel + j];
262
273
  }
263
274
  }
264
- for (let t = 0; t < T; t++)
265
- dX[t] = dLayerIn[t];
266
275
  }
267
- // Embedding lookup: dX at layer-0 input flows into the row for token_t.
268
276
  for (let t = 0; t < T; t++) {
277
+ const { dx, dgain } = rmsNormBackward(dNormedConv[t], lc.layerIn[t], lc.rmsConv[t], nConv);
278
+ for (let c = 0; c < dModel; c++) {
279
+ gNConv[c] = gNConv[c] + dgain[c];
280
+ dLayerIn[t][c] = dLayerIn[t][c] + dx[c];
281
+ }
282
+ }
283
+ return dLayerIn;
284
+ }
285
+ /** Embedding lookup: dL/d(layer-0 input) flows into the row for each token. */
286
+ _embedBackward(tokens, dIn) {
287
+ const { dModel } = this.config;
288
+ for (let t = 0; t < tokens.length; t++) {
269
289
  const off = tokens[t] * dModel;
270
290
  for (let c = 0; c < dModel; c++)
271
- this.gEmb[off + c] = this.gEmb[off + c] + dX[t][c];
291
+ this.gEmb[off + c] = this.gEmb[off + c] + dIn[t][c];
272
292
  }
293
+ }
294
+ /**
295
+ * Next-token cross-entropy over the sequence (predict tokens[t+1] from
296
+ * position t), accumulating exact gradients. Returns the mean loss. Call
297
+ * {@link zeroGrad} before and an optimiser step after.
298
+ */
299
+ lossAndBackward(tokens) {
300
+ if (tokens.length < 2)
301
+ return 0;
302
+ const { logits, cache } = this.forward(tokens);
303
+ const { loss, dX } = this._headBackward(tokens, logits, cache.finalX);
304
+ let d = dX;
305
+ for (let l = this.config.numLayers - 1; l >= 0; l--)
306
+ d = this._backwardLayer(l, d, cache.layers[l]);
307
+ this._embedBackward(tokens, d);
308
+ return loss;
309
+ }
310
+ /**
311
+ * Activation-checkpointed backward — numerically identical gradients to
312
+ * {@link lossAndBackward}, but retains only the per-LAYER inputs during the
313
+ * forward instead of every layer's full activation cache. Each layer's
314
+ * activations are RECOMPUTED (a cheap extra forward) when its backward runs,
315
+ * so peak activation memory is one layer's cache, not all of them. This is the
316
+ * memory-for-compute trade the cookbook pairs with FSDP to fit longer
317
+ * sequences / bigger models on a constrained device.
318
+ */
319
+ lossAndBackwardCheckpointed(tokens) {
320
+ if (tokens.length < 2)
321
+ return 0;
322
+ let x = this._embed(tokens);
323
+ const layerInputs = [];
324
+ for (let l = 0; l < this.config.numLayers; l++) {
325
+ layerInputs.push(x); // keep only the input; drop the cache
326
+ x = this._forwardLayer(l, x).afterMoe;
327
+ }
328
+ const logits = this._head(x);
329
+ const { loss, dX } = this._headBackward(tokens, logits, x);
330
+ let d = dX;
331
+ for (let l = this.config.numLayers - 1; l >= 0; l--) {
332
+ const { cache } = this._forwardLayer(l, layerInputs[l]); // recompute this layer's activations
333
+ d = this._backwardLayer(l, d, cache);
334
+ }
335
+ this._embedBackward(tokens, d);
273
336
  return loss;
274
337
  }
275
338
  // ── Generation ───────────────────────────────────────────────────────────────
@@ -437,26 +500,95 @@ export class EvermindLMTrainer {
437
500
  model;
438
501
  opts;
439
502
  adam;
503
+ scaler;
504
+ grads;
440
505
  constructor(model, opts = {}) {
441
506
  this.model = model;
442
507
  this.opts = opts;
443
508
  this.adam = new AdamW(model, opts);
509
+ this.grads = model.gradients();
510
+ this.scaler = opts.mixedPrecision
511
+ ? new DynamicLossScaler(typeof opts.mixedPrecision === "object" ? opts.mixedPrecision : {})
512
+ : null;
513
+ }
514
+ /** The dynamic loss scaler (mixed-precision mode only) — exposes scale/overflow stats. */
515
+ get lossScaler() {
516
+ return this.scaler;
517
+ }
518
+ _backward(seq) {
519
+ return this.opts.checkpoint ? this.model.lossAndBackwardCheckpointed(seq) : this.model.lossAndBackward(seq);
520
+ }
521
+ /** Divide accumulated gradients by `n` (accumulation averaging), in place. */
522
+ _scaleGrads(n) {
523
+ if (n === 1)
524
+ return;
525
+ const inv = 1 / n;
526
+ for (const g of this.grads)
527
+ for (let i = 0; i < g.data.length; i++)
528
+ g.data[i] = g.data[i] * inv;
529
+ }
530
+ /** Round accumulated gradients to fp16 precision (mixed-precision simulation), in place. */
531
+ _fp16Grads() {
532
+ for (const g of this.grads)
533
+ for (let i = 0; i < g.data.length; i++)
534
+ g.data[i] = roundFp16(g.data[i]);
535
+ }
536
+ /** Multiply accumulated gradients by `s`, in place. */
537
+ _mulGrads(s) {
538
+ if (s === 1)
539
+ return;
540
+ for (const g of this.grads)
541
+ for (let i = 0; i < g.data.length; i++)
542
+ g.data[i] = g.data[i] * s;
444
543
  }
445
544
  /** Train on a set of token sequences; returns per-epoch mean loss. */
446
545
  fit(sequences) {
447
546
  const epochs = this.opts.epochs ?? 1;
547
+ const accum = Math.max(1, this.opts.accumSteps ?? 1);
448
548
  const history = [];
449
549
  for (let e = 0; e < epochs; e++) {
450
550
  let total = 0;
451
551
  let n = 0;
552
+ let pending = 0;
553
+ // The scale is fixed for the duration of one accumulation window, so scale-up
554
+ // (per sequence) and unscale (at flush) always use the SAME factor even as the
555
+ // controller grows/backs it off between windows.
556
+ let windowScale = this.scaler ? this.scaler.scale : 1;
557
+ this.model.zeroGrad();
558
+ const flush = () => {
559
+ if (pending === 0)
560
+ return;
561
+ this._scaleGrads(pending); // accumulation averaging
562
+ if (this.scaler) {
563
+ this._fp16Grads(); // half-precision gradients
564
+ const overflow = this.scaler.check(this.grads);
565
+ if (this.scaler.update(overflow)) {
566
+ this._mulGrads(1 / windowScale); // unscale by the SAME factor we scaled by
567
+ this.adam.step();
568
+ } // else: skip the step, scale has backed off
569
+ windowScale = this.scaler.scale; // next window uses the updated scale
570
+ }
571
+ else {
572
+ this.adam.step();
573
+ }
574
+ this.model.zeroGrad();
575
+ pending = 0;
576
+ };
452
577
  for (const seq of sequences) {
453
578
  if (seq.length < 2)
454
579
  continue;
455
- this.model.zeroGrad();
456
- total += this.model.lossAndBackward(seq);
457
- this.adam.step();
580
+ if (pending === 0 && this.scaler)
581
+ windowScale = this.scaler.scale;
582
+ const loss = this._backward(seq);
583
+ // Emulate scaling the loss before backward: the gradient is linear in the
584
+ // loss, so scaling the just-produced gradients by the loss scale is equivalent.
585
+ this._mulGrads(windowScale);
586
+ total += loss;
458
587
  n++;
588
+ if (++pending >= accum)
589
+ flush();
459
590
  }
591
+ flush();
460
592
  history.push(n > 0 ? total / n : 0);
461
593
  }
462
594
  return history;