fix: give each translation worker its share of the CPUs

Translate runs NumCPU/2 sentence workers and each of them started
GOMAXPROCS goroutines for the logit projection of every step: on 16
cores, 128 goroutines fighting for 16 cores. The model now caps the
goroutines of a step (SetThreads, called once after Load) so that
workers x threads stays within the machine.
This commit is contained in:
Kharec committed 2026-10-15 17:53:38 +02:00
1 parent d6ee57130f
commit 82bfef8cba
4 files changed
+26 -7

No files matched your search

+3
View File
@@ -212,6 +212,9 @@ func (m *Model) decodeStep(b *beam, crossK, crossV [][]float32, srcLen int, s *d
func (m *Model) logProbs(hidden, out [][]float32) [][]float32 {
d, vocab := m.cfg.DModel, m.cfg.VocabSize
workers := runtime.GOMAXPROCS(0)
if m.threads > 0 {
workers = min(workers, m.threads)
}
chunk := (vocab + workers - 1) / workers
var wg sync.WaitGroup
for lo := 0; lo < vocab; lo += chunk {
+11 -2
View File
@@ -60,10 +60,12 @@ type decoderLayer struct {
finalLayerNorm layerNorm
}
// Model is a loaded translation model. It is read-only once loaded:
// Translate may be called from several goroutines at once.
// Model is a loaded translation model. It is read-only once loaded (but
// for SetThreads, right after Load): Translate may be called from several
// goroutines at once.
type Model struct {
cfg config
threads int // goroutines a decoding step may start (0: one per CPU)
embed []float32 // vocab × d, shared by encoder, decoder and output
encPos []float32 // positions × d
decPos []float32
@@ -221,6 +223,13 @@ func (m *Model) bind(t map[string]tensor) error {
return missing
}
// SetThreads caps how many goroutines each decoding step starts (0: one
// per CPU, the default). When several sentences are translated at once,
// give each its share of the CPUs, or every step of every sentence starts
// a team of goroutines and they oversubscribe the machine. It must be
// called before the first Translate.
func (m *Model) SetThreads(n int) { m.threads = n }
// maxSourceTokens leaves room for the end-of-sentence token.
func (m *Model) maxSourceTokens() int { return m.cfg.MaxPositions - 1 }