librispeech deepspeech validation loss differs between pytorch and jax, but wer aligns

[wyfEmma]

Provide a general description of the issue
validation loss:

wer:

[wyfEmma]

this is how jax submission calculates ctc_loss :

submission_runner calls eval_model

algorithmic-efficiency/submission_runner.py

Line 465 in 5849117

latest_eval_result = workload.eval_model(

eval_model calls eval_model_on_split for the validation set

algorithmic-efficiency/algoperf/spec.py

Line 347 in 5849117

validation_metrics = self._eval_model_on_split(

then it calls eval steps, which calls _eval_step:

algorithmic-efficiency/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py

Line 362 in 5849117

def eval_step(

[wyfEmma]

in jax workload, to get the number of examples across all devices, it sums the loss['n_valid_examples'] across all devices.
which means it assumes that loss['n_valid_examples'] represents examples on one GPU and summing num of examples on all devices gives us the number of valid examples on all GPUs

algorithmic-efficiency/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py

Line 357 in 5849117

'n_valid_examples': jnp.zeros((len(jax.devices()), 1))

[wyfEmma]

however, printing the loss['n_valid_examples'], it gives us the shape of 256

, which matches the eval_batch_size of librispeech

algorithmic-efficiency/algoperf/workloads/librispeech_conformer/workload.py

Line 70 in 5849117

def eval_batch_size(self) -> int:

this means that the loss['n_valid_examples'] already hold data across all devices, multiplying it by device_count upscaled the number of examples by 4.

when ctc_loss is calculated, in jax the denominator is 4 times larger, so the loss is 4 times smaller, this explains the difference we see in the ctc_loss

[wyfEmma]

I think the reason that loss[n_valid_example] contains the data cross all device is because we are now using jit instead of pmap

the fix: remove the replication by 4 here

algorithmic-efficiency/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py

Line 357 in 5849117

'n_valid_examples': jnp.zeros((len(jax.devices()), 1))

[wyfEmma]

after the fix ctc loss aligns:

step graph with CTC_loss:

time plot: