Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix bugs in workload variants #658

Merged
merged 3 commits into from
Feb 28, 2024
+4 −3
Merged

Fix bugs in workload variants #658

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8ce8772
Fix WMT jax config for decoding
runame Feb 28, 2024
74124d5
Fix deepspeech model_state when batchnorm is not used
runame Feb 28, 2024
a4dff63
Fix lint
runame Feb 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion algorithmic_efficiency/workloads/librispeech_conformer/librispeech_jax/workload.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -327,7 +327,7 @@ def _eval_model_on_split(self,
global_step: int = 0) -> Dict[str, float]:
"""Run a full evaluation of the model."""
del global_step
if model_state is not None:
if model_state is not None and len(model_state) > 0:
# Sync batch statistics across replicas before evaluating.
model_state = self.sync_batch_stats(model_state)

Expand Down
3 changes: 2 additions & 1 deletion algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_jax/workload.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,8 @@ def init_model_fn(
variables = model_init_fn({'params': params_rng, 'dropout': dropout_rng},
*fake_input_batch)

model_state = variables['batch_stats']
model_state = variables[
'batch_stats'] if not self.layernorm_everywhere else {}
params = variables['params']
self._param_shapes = param_utils.jax_param_shapes(params)
self._param_types = param_utils.jax_param_types(self._param_shapes)
Expand Down
2 changes: 1 addition & 1 deletion algorithmic_efficiency/workloads/wmt/wmt_jax/workload.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ def predict_step(self,
max_decode_len: int,
beam_size: int = 4) -> spec.Tensor:
"""Predict translation with fast decoding beam search on a batch."""
config = models.TransformerConfig(deterministic=True, decode=True)
config = replace(self._eval_model.config, decode=True)
# Prepare transformer fast-decoder call for beam search: for beam search, we
# need to set up our decoder model to handle a batch size equal to
# batch_size * beam_size, where each batch item's data is expanded in-place
Expand Down
Loading