[Hotfix/randomizer]Fix Randomizer error on CPU #5265
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
馃搶 Checklist before creating the PR
[doc/gemini/tensor/...]: A concise description
馃毃 Issue number
馃摑 What does this PR do?
Previously, Randomizer doesn't support cpu due to the
_dispatch_device_func
, which causes OpenMoE unable to initialize parameters on CPU. When converting a 34B OpenMoE Checkpoint, I found this inconvenient because I can only convert a 166GB checkpoint in CPU RAM (I don't have such a huge GPU memory for that).The bug is raise from: colossalai/moe/experts.py#L88, where
set_rng_state
andget_rng_state
need to be executed during parameters initialization. I found on cpu, you can usetorch. set/get_rng_state
to manipulate the rng_state. So I modify the_dispatch_device_func
, allow to dispatch supported functions likeget/set_rng_state
to cpu, and still raise error if some functions are not supported on cpu device, such asget/set_rng_state_all
.In addition to fixing bugs, I also added some conditional statements to make the code more robust. The modified code is compatible with original behaviors. After modification, you can run code below to verify the effectiveness
馃挜 Checklist before requesting a review
猸愶笍 Do you enjoy contributing to Colossal-AI?
Tell us more if you don't enjoy contributing to Colossal-AI.