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.

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

Improve CI speed and resolve issues of run_quantization_check #1682

Merged
merged 1 commit into from
Jul 8, 2024

Conversation

james77777778
Copy link
Collaborator

@james77777778 james77777778 commented Jul 4, 2024

I have analyzed the time cost in run_quantization_check

Test command: CUDA_VISIBLE_DEVICES= KERAS_BACKEND=jax pytest keras_nlp/src/ -k backbone_basics

time cost
run_quantization_check=False 45.39s
run_quantization_check=True without calling quantize 47.45s
run_quantization_check=True with calling quantize 269s
run_quantization_check=True with calling quantize + saving 275s

Obviously, the bottleneck is the underlying computation when calling Model.quantize.

Here, I propose an improvement by pre-configuring DTypePolicyMap and using it to instantiate the model to avoid quantization-related computation.
This should improve the speed of CI.

Some minor bugs have been spotted and resolved too.

  • Backbone.get_config() should consider that self.dtype_policy is already a DTypePolicyMap
  • The name of self.embeddings_layer_norm in BloomBackbone has been changed to avoid a duplicated name that breaks DTypePolicyMap.
  • OPTBackbone.get_config missed super()
  • XLNetBackbone failed to pass run_quantization_check. (Will try to fix it in the future)

EDITED:
The CI is much faster now. (JAX: ~27mins -> 18mins)

Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

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

This looks great! Thank you!

@mattdangerw mattdangerw added the kokoro:force-run Runs Tests on GPU label Jul 8, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 8, 2024
@mattdangerw mattdangerw merged commit f9faaf1 into keras-team:master Jul 8, 2024
8 checks passed
@james77777778 james77777778 deleted the improve-test-speed branch July 9, 2024 00:07
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.

3 participants