Skip to content

[Flax] Bugfixes in run_mlm_flax.py#9127

Closed
TevenLeScao wants to merge 4 commits intohuggingface:masterfrom
TevenLeScao:run_mlm_flax_bugfixes
Closed

[Flax] Bugfixes in run_mlm_flax.py#9127
TevenLeScao wants to merge 4 commits intohuggingface:masterfrom
TevenLeScao:run_mlm_flax_bugfixes

Conversation

@TevenLeScao
Copy link
Contributor

What does this PR do?

This PR fixes a few bugs I have observed when using run_mlm_flax.py:

  • As discussed with @mfuntowicz , jnp.split is a lot slower than np.split on the first iteration, outright hanging in my tests on simplewiki (~20MB). As this operation doesn't need to be traced. we can use np.split instead.
  • When using a HF datasets, the text column was also passed to the model as input, causing a bug. The PR removes the text column in dataset.map to avoid this.
  • Finally, using warmup_steps = 0 (as is default) causes the Flax optimizer to output NaNs. We use 1 as a minimum value for the same warmup-less behaviour.

Copy link
Collaborator

@sgugger sgugger left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks!

).create(model.params)

# Create learning rate scheduler
# warmup_steps = 0 causes the Flax optimizer to return NaNs; warmup_steps = 1 is functionally equivalent.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could a do a logger.warning(...) if training_args.warmup_steps happens to be 0.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just added !

@github-actions
Copy link
Contributor

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

@github-actions github-actions bot closed this Apr 25, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants