Handle nested fp8 autocasts

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True

[ksivaman]

Fixes the case of handling nested fp8_autocasts as well as moving from fp8 → high precision training. See comment for example.

[ksivaman]

/te-ci

[timmoon10]

If I understand correctly, the error happened because TransformerEngineBaseModule.fp8 didn't have logic to be set to False, so it could have an incorrect value if we alternate between FP8 and non-FP8. This led us down the FP8 code path at unexpected times, e.g.:

TransformerEngine/transformer_engine/pytorch/module.py

Line 616 in 36e0ac5

ifself.fp8andself.training:

This fix looks reasonable to me. I wonder if we'll hit an edge case in is_first_fp8_module though.

[timmoon10]

Could we run into a bug if the user does something like:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

I think the existing implementation also has a bad edge case:

withfp8_autocast():
withfp8_autocast():
passwithfp8_autocast():
module1(x) # is_first_fp8_module() == False

I think the right answer is to remove _IS_FIRST_FP8_MODULE entirely from the finally clause, but I'm not confident if there are other interactions I'm missing.

[timmoon10]

After discussion with @ksivaman, we currently only support nested fp8_autocast s in the case we want to disable FP8 in certain modules:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
module2() # is_first_fp8_module() == True

This PR should handle this correctly. We should expect bad behavior with other nesting patterns:

withfp8_autocast():
withfp8_autocast():
module1() # is_first_fp8_module() == Truewithfp8_autocast():
module2() # is_first_fp8_module() == True

This is somewhat worrying if we want to reenable FP8 in some region:

withfp8_autocast():
withfp8_autocast(enabled=False):
module1()
withfp8_autocast(enabled=True)
module2() # is_first_fp8_module() == Truemodule3() # is_first_fp8_module() == True