Clip convert error: Dynamic value of min/max is not implemented about onnx2torch

[ongiaf]

The full onnx model can be download from here:

1. https://get.ecmwf.int/repository/test-data/ai-models/fuxi/short.onnx
2. ONNX External Data: https://get.ecmwf.int/repository/test-data/ai-models/fuxi/short

[dsuhoi]

@ongiaf Do you have any success decoding FuXi (I've been fine-tuning this model for a long time)? I recommend paying attention to this solution

[ongiaf]

Thanks, it's excellent work.
And with some dirty work, Fuxi can successfully run on PyTorch with Onnx2Torch. In Onnx2Torch, problems are mainly about LayerNormalization and Clip.

[dsuhoi]

@ongiaf Did you manage to run FuXi with the current weights for the fine-tuning process (I am currently thinking about how to complete the work on the model on a 1-hour grid and thought about freezing some layers except the U-transformer.) ?

[juanqiu1]

Thanks, it's excellent work. And with some dirty work, Fuxi can successfully run on PyTorch with Onnx2Torch. In Onnx2Torch, problems are mainly about LayerNormalization and Clip.

Thank you for posting your changes about Clip. Could you also suggest how to fix LayerNormalization? Looks like the converted model has issue with torch.layer_norm call.

[dsuhoi]

@juanqiu1 In order for this to work with FuXi, you will need to change the onnx2torch/node_converters/layer_norm.py parameter to [1536]:

@add_converter(operation_type='LayerNormalization', version=17)
def _(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
    node_attributes = node.attributes

    axis = node_attributes.get('axis', AXIS_DEFAULT_VALUE)
    epsilon = node_attributes.get('epsilon', EPSILON_DEFAULT_VALUE)

    if all(value_name in graph.initializers for value_name in node.input_values[1:]):
        input_value_info = graph.value_info[node.input_values[0]]
        input_shape = get_shape_from_value_info(input_value_info)
        torch_module = nn.LayerNorm(
            normalized_shape=(1536), # input_shape[axis:], (this block!)
            eps=epsilon,
            elementwise_affine=True,
        )

[juanqiu1]

@dsuhoi Thank you for hinting, there are a couple of other easy fixes (typing, etc).
Did you manage to run FuXi with the current weights for the fine-tuning process
Do you have any progress on that? After conversion, I loaded model into pytorch but even on A100 with FSDP enabled via accelerate. I still get CUDA out-of-memory error.

[dsuhoi]

@juanqiu1 Yes, I managed to start the learning process by highlighting the named_parameters() part within the last dozen UTransformer (this was enough for fine-tunning).

I used Nvidia A100 (40GB).