Code

Add JSD Loss for Distillation#425

Code

Merged

Tcc0403 merged 11 commits intomainlinkedin/Liger-Kernel:mainfrom austin362667/chunked_compiled_jsd_losslinkedin/Liger-Kernel:austin362667/chunked_compiled_jsd_lossCopy head branch name to clipboardJan 30, 2025

+781-7

Commits

Commits on Jan 23, 2025

• Add chunked and compiled JSD loss

  Show description for d2d973f
  austin362667

  committed
  d2d973f

  Copy full SHA for d2d973f

• Add tests for JSD loss

  Show description for 8521421
  austin362667

  committed
  8521421

  Copy full SHA for 8521421

• Add benchmarks for JSD loss

  Show description for 9218cf0
  austin362667

  committed
  9218cf0

  Copy full SHA for 9218cf0

• Set JSD loss default chunk_size to 2

  Show description for c53a46c
  austin362667

  committed
  c53a46c

  Copy full SHA for c53a46c

• refactor temperature scaling and fix chunked distillation loss norm reduction

  Show description for af3f83a
  austin362667

  committed
  af3f83a

  Copy full SHA for af3f83a

• adjust tol

  Show description for d533471
  austin362667

  committed
  d533471

  Copy full SHA for d533471

• Use log space and add beta for jsd loss

  Show description for 46fb935
  austin362667

  committed
  46fb935

  Copy full SHA for 46fb935

• nit

  Show description for f047070
  austin362667

  committed
  f047070

  Copy full SHA for f047070

• expose beta

  Show description for 50022c0
  austin362667

  committed
  50022c0

  Copy full SHA for 50022c0

• align jsd with paper & twist tests beta

  Show description for 7268f64
  austin362667

  committed
  7268f64

  Copy full SHA for 7268f64

Commits on Jan 29, 2025

• polish docstring

  Show description for 872169f
  austin362667

  committed
  872169f

  Copy full SHA for 872169f