experimental support to LOMO optimizer

[zhenqincn]

Add support to LOMO optimizer for LLM Full Parameter Fine-tuning for Large Language Models with Limited Resources

[CLAassistant]

All committers have signed the CLA.

[rayrayraykk]

Since FS has TF backend, we should use try-catch to avoid error:

try:
    import torch
except ImportError:
    torch=None

[zhenqincn]

Currently, LOMO has no support to TF. A try-catch statement is added following this suggestion as:

try:
    import torch
except ImportError:
    torch=None
    raise ImportError('Currently, LOMO optimizer is only implemented with `pytorch`')

[rayrayraykk]

Please don't raise ImportError or anything in header (do it in __init__ or somewhere), since tf backend user might run into error since this file will be imported.

[zhenqincn]

OK, in the latest commit, this error raising has beem moved in __init__ in the corresponding optimizer.

try:
    import torch
except ImportError:
    torch=None

class LOMO(Optimizer):
    """
    an optimizer for LOMOTrainer
    """

    def __init__(self, model, lr=1e-3, clip_grad_norm=None, clip_grad_value=None):
        if torch is None:
            raise ImportError('Currently, LOMO optimizer is only implemented with `pytorch`')
        self.model = model

[rayrayraykk]

use logger

[zhenqincn]

This is an omission during development, which has been removed in a new commit.

[rayrayraykk]

this check should be in _hook_on_fit_start_init

[zhenqincn]

This check has been moved in _hook_on_fit_start_init

    def _hook_on_fit_start_init(self, ctx):
        ret = super()._hook_on_fit_start_init(ctx)
        if not isinstance(ctx.optimizer, LOMO):
            raise AttributeError(f'"lomo" must be set as the type of ',
                                 f'`train.optimizer` if the trainer is LOMOTrainer')
        return ret

[rayrayraykk]

Since the train and eval will use the same hook_func, we should add an if-else when eval only needs one forward, right?

[zhenqincn]

Sure, one additional check has been added as following:

if ctx.cur_mode in [MODE.TRAIN, MODE.FINETUNE] \
    and (
            not ctx.skip_this_batch 
            and ctx.optimizer.clip_grad_norm is not None 
            and ctx.optimizer.clip_grad_norm > 0
        ):

[rayrayraykk]

Another minor issue is that since the FS-LLM is publicly available now, the dev/llm branch is deprecated. And you can change the target branch of this PR to llm, thx!

[zhenqincn]

All suggestions provided above have been adopted. A review is re-requested. Many thanks.

[zhenqincn]

In the latest commit, the code has been formatted with pre-commit checks passed.

[rayrayraykk]

torch.Optimizer should be in try-catch.

[zhenqincn]

Many thanks. It has been solved.

[yxdyc]

Line 87 ~ 115：
There are too many judgement branches and they are too deeply nested.I suggest changing it to a series of self-explanatory boolean variables to increase readability

[zhenqincn]

Thanks for this suggestion. The mentioned lines have been reformatted with code comments added in the latest commit.

[yxdyc]

perhaps add some reference to make the checking reasonable, e.g., stage3_code

[zhenqincn]

Many thanks. This reference has been added in the latest commit.