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 -1
View File
@@ -23,7 +23,9 @@ version is in the `VERSION` file.
replacing the link.
- Translation is faster and uses much less memory: the beam search no longer copies
every beam's decoder state for every candidate at every step, and reuses its
work buffers instead of allocating new ones per layer and step.
work buffers instead of allocating new ones per layer and step. Each sentence
now uses its share of the processors, so translating a script no longer starts
more workers than the machine has.
## [1.4] - 2026-10-15
+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 }
+9 -4
View File
@@ -232,10 +232,8 @@ func (a *App) Translate(doc screenplay.Document, elements []int, titlePage bool)
res, err := translate.Document(ctx, &doc, model, lexicon, translate.Options{
Elements: elements,
TitlePage: titlePage,
// The model splits its largest product across every CPU already:
// half as many sentences at once keeps the machine responsive.
Workers: max(1, runtime.NumCPU()/2),
Progress: func(done, total int) { emit(int64(done), int64(total)) },
Workers: translateWorkers(),
Progress: func(done, total int) { emit(int64(done), int64(total)) },
})
if errors.Is(err, context.Canceled) {
return nil, nil
@@ -246,6 +244,10 @@ func (a *App) Translate(doc screenplay.Document, elements []int, titlePage bool)
return &res, nil
}
// translateWorkers is how many sentences are translated at once: half the
// CPUs, and each one splits its decoding steps across its share of them.
func translateWorkers() int { return max(1, runtime.NumCPU()/2) }
// translator returns the model and the dictionaries that help tell French
// from English, loading them the first time. The load happens without
// a.tr.mu: reading 300 MB takes seconds, and CancelTranslation must stay
@@ -271,6 +273,9 @@ func (a *App) translator(ctx context.Context) (*marian.Model, translate.Lexicon,
}
return nil, nil, fmt.Errorf("loading the translation model: %w", err)
}
// Every sentence worker gets an equal share of the CPUs for its
// decoding steps, so workers × threads never oversubscribe the machine.
m.SetThreads(max(1, runtime.NumCPU()/translateWorkers()))
lexicon = loadLexicon()
a.tr.mu.Lock()
if ctx.Err() == nil { // cancelled while loading: cache nothing