@kurisusnowdeng The checks are not passed

Hi @kurisusnowdeng , I have reviewed your code and some issues remain.

benchmark/cifar/configs/vit_1d.py

1. the correct field name for gradient_clipping is clip_grad_norm. This applies to all other config files.

colossalai/nn/init.py

1. seems fan_in and fan_out can be directly obtained from the weight shape, they should be hidden from the user.
2. fan_in should not be set to None as default because it cannot be default in the assert statement in lecun initialization.

def lecun_uniform_():
    # adapted from jax.nn.initializers
    def initializer(tensor: Tensor, fan_in: int = None, fan_out: int = None):
        assert fan_in is not None, 'Fan_in is not provided.'

        var = 1.0 / fan_in
        bound = math.sqrt(3 * var)
        return nn.init.uniform_(tensor, -bound, bound)

    return initializer


def lecun_normal_():
    # adapted from jax.nn.initializers
    def initializer(tensor: Tensor, fan_in: int = None, fan_out: int = None):
        assert fan_in is not None, 'Fan_in is not provided.'

        std = math.sqrt(1.0 / fan_in)
        return nn.init.trunc_normal_(tensor, std=std / .87962566103423978)

    return initializer

colossalai/nn/layer/__init__.py

1. from line 6 to line 24, why are all the modules imported?
2. Are there reasons for setting the default initializer? e.g. kaiming for weight and xavier for bias. The user is not aware of our setting and my suggestion would be to fall back to default torch initialization.
3. I would suggest putting the unified layers into a single file (e.g. unified_layer.py) instead of __init__.py as __init__.py serves for intialization/configuration purpose.

colossalai/nn/layer/parallel_1d/_utils.py

1. the utils functions here are for general use, not specific for 1D, I don’t think they should be put under 1D utils.

colossalai/nn/metric/accuracy_2d.py

1. accuracy is still using autograd function and this applies to other metrics as well.

Fixed. Please check again. @FrankLeeeee