[RFC] [Relay] Automatic Mixed Precision Pass

[AndrewZhaoLuo]

Relevant Links:

https://discuss.tvm.apache.org/t/rfc-relay-fp32-fp16-model-support/9994

• Old discussion before the new RFC process was rolled out

apache/tvm#8069

• Initial PR

cc @hogepodge@mbrookhart@anijain2305@masahi

Link to tracking issue: apache/tvm#8296

[comaniac]

Thanks for the RFC. I have two questions:

1. How to mark/set the color (i.e., attribute) of every operator?

2. It seems to me that if we register a casting checker instead of just a label (color), then we can simplify the algorithm a lot. Taking the case A(green) - B(gray) - C(green) as an example, if we could register a casting rule of B as follows, then we just need one traverse to know if we need cast around B:

   def amp_B(expr, args):
   a = args[0]
   if (a.dtype is float16):
   return fp16
   return fp32
   

   After all, we only need the previous nodes to determine 1) whether to use FP16 implementation, and 2) whether to insert casts. It seems to me that this pass is similar to the layout conversion pass, which uses one traverse to finish everything, so it might be possible for AMP too.

[AndrewZhaoLuo]

Thanks for the RFC. I have two questions:

1. How to mark/set the color (i.e., attribute) of every operator?

2. It seems to me that if we register a casting checker instead of just a label (color), then we can simplify the algorithm a lot. Taking the case A(green) - B(gray) - C(green) as an example, if we could register a casting rule of B as follows, then we just need one traverse to know if we need cast around B:

   def amp_B(expr, args):
   a = args[0]
   if (a.dtype is float16):
   return fp16
   return fp32
   

   After all, we only need the previous nodes to determine 1) whether to use FP16 implementation, and 2) whether to insert casts. It seems to me that this pass is similar to the layout conversion pass, which uses one traverse to finish everything, so it might be possible for AMP too.

Yep that is correct it is very similar to the layout conversion pass. This RFC has an initial PR here: apache/tvm#8069.

To answer your questions:

1. src/relay/transforms/fp32_to_fp16.h -- DefaultFP16Colorer is the default way. But the only thing we need is a callable with type CallNode*(Color). So you could write your own colorer that does arbitrary stuff when only looking at a single node at a time.

2. This is functionally what is done in the PR I link. It's one pass.

[comaniac]

Thanks for the answers. I'll review the PR to get more implementation details.
One more question regarding the extensibility: can this be extended easily to support bfloat16?

[AndrewZhaoLuo]

Thanks for the answers. I'll review the PR to get more implementation details.
One more question regarding the extensibility: can this be extended easily to support bfloat16?

It should be trivial (hope I don't eat my words). I'm not 100% sure of the support for bfloat16 in current relay ops however.

[AndrewZhaoLuo]

I don't know Chris Sullivan's github handle so if someone could cc him too that would be great.

[tmoreau89]

CCing @csullivan

[comaniac]

Thanks for the answers. I'll review the PR to get more implementation details.
One more question regarding the extensibility: can this be extended easily to support bfloat16?

It should be trivial (hope I don't eat my words). I'm not 100% sure of the support for bfloat16 in current relay ops however.

TVM has limited bfloat16 support now but it's on the way, so it would be better for this RFC to also consider this case, even the initial version may not cover it.

[AndrewZhaoLuo]

So the associated PR is getting closer to a mergeable state. Is this RFC ready for more comments?

[tqchen]

cc @comaniac would be great if you can help shepherd this RFC

[comaniac]

The concept and algorithm look good to me, but it would be better to provide more implementation/design details.

[comaniac]

The answer to this question should come with a discussion of existing mechanisms used by other frameworks, such as XLA and PyTorch.

[AndrewZhaoLuo]

Done. Please let me know if this is sufficient. Don't have the best background on some of this stuff.

[AndrewZhaoLuo]

Thanks for driving this review @comaniac. I'll get to this later in the week.

[AndrewZhaoLuo]

Going to get to this tomorrow 😬. Promise 🤞

[comaniac]

btw, according to #17, please update the RFC number on the file name to align with this PR number.

[comaniac]

Took a quick pass to the updated RFC. I think it's almost ready to merge as long as the last 3 comments are resolved.

[AndrewZhaoLuo]

PTAL @comaniac

[AndrewZhaoLuo]

@comaniac, I'll be talking about this at the TVM community meeting tomorrow so put off merging until after.

[comaniac]

LGTM. Will merge after the community meeting if there's no objection.

[AndrewZhaoLuo]

If there is not other objections, this will be merged on monday.

[comaniac]

Thanks @AndrewZhaoLuo

[MeJerry215]

@AndrewZhaoLuo will it remove cast weight to float16 from graph? and make weight as float16 when build lib.
in my opinion, it will reduce the bandwidth.

[masahi]

@MeJerry215 Yes, casting of weight to fp16 is done at compile time by FoldConstant pass, so weights will be in fp16 at deploy time.