【训练营】学习率调度器实现 #113
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,186 @@ | ||
| #pragma once | ||
|
|
||
| #include <cmath> | ||
| #include <cstdint> | ||
| #include <functional> | ||
| #include <memory> | ||
| #include <string> | ||
| #include <unordered_map> | ||
| #include <variant> | ||
| #include <vector> | ||
|
|
||
| namespace infini_train { | ||
|
|
||
| class Optimizer; | ||
|
|
||
| using StateValue = std::variant<int64_t, float, double, std::string, std::vector<float>>; | ||
| using StateDict = std::unordered_map<std::string, StateValue>; | ||
|
|
||
| struct LRSchedulerConfig { | ||
| std::string type = "none"; | ||
| // ConstantLR | ||
| float constant_factor = 1.0f / 3.0f; | ||
| int constant_total_iters = 5; | ||
| // StepLR | ||
| int64_t step_size = 10; | ||
| float step_gamma = 0.1f; | ||
| // LinearLR | ||
| float linear_start_factor = 1.0f / 3.0f; | ||
| float linear_end_factor = 1.0f; | ||
| int linear_total_iters = 5; | ||
| // LambdaLR | ||
| std::function<float(int64_t)> lambda_fn = nullptr; | ||
| // SequentialLR | ||
| std::vector<LRSchedulerConfig> sequential_configs; | ||
| std::vector<int64_t> sequential_milestones; | ||
| // ChainedScheduler | ||
| std::vector<LRSchedulerConfig> chained_configs; | ||
| // warmup | ||
| int64_t warmup_steps = 0; | ||
| float warmup_start_factor = 1.0f / 3.0f; | ||
| float warmup_end_factor = 1.0f; | ||
| }; | ||
|
|
||
| class LRScheduler { | ||
| public: | ||
| template <typename T, typename... Args> static std::shared_ptr<T> Create(Args &&...args) { | ||
| auto scheduler = std::make_shared<T>(std::forward<Args>(args)...); | ||
| scheduler->InitialStep(); | ||
| return scheduler; | ||
| } | ||
|
|
||
| explicit LRScheduler(std::shared_ptr<Optimizer> optimizer, int64_t last_step = -1); | ||
| virtual ~LRScheduler() = default; | ||
|
|
||
| LRScheduler(const LRScheduler &) = delete; | ||
| LRScheduler &operator=(const LRScheduler &) = delete; | ||
|
|
||
| virtual void Step(); | ||
| virtual void Step(int64_t epoch); | ||
| virtual void InitialStep(); | ||
|
|
||
| float GetLR() const; | ||
| float BaseLR() const; | ||
| int64_t LastStep() const; | ||
|
|
||
| void ResetStep(int64_t step = -1); | ||
| virtual StateDict State() const; | ||
| virtual void LoadState(const StateDict &state); | ||
|
|
||
| protected: | ||
| virtual float GetClosedFormLR() const = 0; | ||
| virtual float GetChainedFormLR() const; | ||
| void ApplyLR(float lr); | ||
|
|
||
| std::shared_ptr<Optimizer> optimizer_; | ||
| int64_t last_step_; | ||
| float current_lr_; | ||
|
||
| float base_lr_; | ||
| bool is_initial_ = false; | ||
| }; | ||
|
|
||
| std::shared_ptr<LRScheduler> CreateLRScheduler(std::shared_ptr<Optimizer> optimizer, const LRSchedulerConfig &config); | ||
|
|
||
| namespace lr_schedulers { | ||
|
|
||
| class ConstantLR : public LRScheduler { | ||
| public: | ||
| ConstantLR(std::shared_ptr<Optimizer> optimizer, float factor = 1.0f / 3.0f, int total_iters = 5, | ||
| int64_t last_step = -1); | ||
| ~ConstantLR() override = default; | ||
|
|
||
| protected: | ||
| float GetChainedFormLR() const override; | ||
| float GetClosedFormLR() const override; | ||
|
|
||
| private: | ||
| const float factor_; | ||
| const int64_t total_iters_; | ||
| }; | ||
|
|
||
| class StepLR : public LRScheduler { | ||
| public: | ||
| StepLR(std::shared_ptr<Optimizer> optimizer, int64_t step_size, float gamma = 0.1f, int64_t last_step = -1); | ||
| ~StepLR() override = default; | ||
|
|
||
| protected: | ||
| float GetChainedFormLR() const override; | ||
| float GetClosedFormLR() const override; | ||
|
|
||
| private: | ||
| const int64_t step_size_; | ||
| const float gamma_; | ||
| }; | ||
|
|
||
| class LinearLR : public LRScheduler { | ||
| public: | ||
| LinearLR(std::shared_ptr<Optimizer> optimizer, float start_factor = 1.0f / 3.0f, float end_factor = 1.0f, | ||
| int64_t total_iters = 5, int64_t last_step = -1); | ||
| ~LinearLR() override = default; | ||
|
|
||
| protected: | ||
| float GetChainedFormLR() const override; | ||
| float GetClosedFormLR() const override; | ||
|
|
||
| private: | ||
| const float start_factor_; | ||
| const float end_factor_; | ||
| const int64_t total_iters_; | ||
| }; | ||
|
|
||
| class LambdaLR : public LRScheduler { | ||
| public: | ||
| using LambdaFunc = std::function<float(int64_t)>; | ||
|
|
||
| LambdaLR(std::shared_ptr<Optimizer> optimizer, LambdaFunc lr_lambda, int64_t last_step = -1); | ||
| ~LambdaLR() override = default; | ||
|
|
||
| protected: | ||
| float GetClosedFormLR() const override; | ||
|
|
||
| private: | ||
| const LambdaFunc lr_lambda_; | ||
| }; | ||
|
|
||
| class SequentialLR : public LRScheduler { | ||
| public: | ||
| SequentialLR(std::shared_ptr<Optimizer> optimizer, std::vector<std::shared_ptr<LRScheduler>> schedulers, | ||
| std::vector<int64_t> milestones, int64_t last_step = -1); | ||
| ~SequentialLR() override = default; | ||
|
|
||
| void Step() override; | ||
| void InitialStep() override; | ||
|
|
||
| StateDict State() const override; | ||
| void LoadState(const StateDict &state) override; | ||
|
|
||
| protected: | ||
| float GetClosedFormLR() const override { return current_lr_; } | ||
|
||
| void UndoChildInitialSteps(); | ||
|
|
||
| private: | ||
| std::vector<std::shared_ptr<LRScheduler>> schedulers_; | ||
| std::vector<int64_t> milestones_; | ||
| }; | ||
|
|
||
| class ChainedScheduler : public LRScheduler { | ||
| public: | ||
| ChainedScheduler(std::shared_ptr<Optimizer> optimizer, std::vector<std::shared_ptr<LRScheduler>> schedulers, | ||
| int64_t last_step = -1); | ||
| ~ChainedScheduler() override = default; | ||
|
|
||
| void Step() override; | ||
| void InitialStep() override; | ||
|
|
||
| StateDict State() const override; | ||
| void LoadState(const StateDict &state) override; | ||
|
|
||
| protected: | ||
| float GetClosedFormLR() const override { return current_lr_; } | ||
|
|
||
| private: | ||
| std::vector<std::shared_ptr<LRScheduler>> schedulers_; | ||
| }; | ||
|
|
||
| } // namespace lr_schedulers | ||
| } // namespace infini_train | ||
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
scheduler 在前面已经 Step 过了，所以这里 GetLR() 语义上是”下一步要用到的 lr“；而我们这里想打印的是每一步实际用到的 lr，所以这里的逻辑需要修改下。llama3 部分的 main.cc 里同理。