Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
Apache License 2.0
1.12k stars 69 forks source link

Support for Python 3.12 #763

Open t-vi opened 1 month ago

t-vi commented 1 month ago

We should look into Python 3.12 somewhat soonish: The bulk of the work required probably can be inferred from the Python dis documentation by searching for 3.12 changes.

The list below is as of #767 which makes the test_interpreter on 3.12 pass and thus enables us to have CI checking the Thunder Python Interpreter.

While the tests pass, clearly, they do not cover the unimplemented opcodes, so we should look into closing the delta for that.

For context, the Ubuntu 24.04 LTS uses Python 3.12.

t-vi commented 1 month ago
In [11]: implemented - tested_instructions
Out[11]: {'LOAD_ASSERTION_ERROR', 'LOAD_METHOD'}

In [12]: {*dis.opmap} - implemented
Out[12]: 
{'CACHE',
 'CALL_INTRINSIC_2',
 'CHECK_EG_MATCH',
 'CLEANUP_THROW',
 'INSTRUMENTED_CALL',
 'INSTRUMENTED_CALL_FUNCTION_EX',
 'INSTRUMENTED_END_FOR',
 'INSTRUMENTED_END_SEND',
 'INSTRUMENTED_FOR_ITER',
 'INSTRUMENTED_INSTRUCTION',
 'INSTRUMENTED_JUMP_BACKWARD',
 'INSTRUMENTED_JUMP_FORWARD',
 'INSTRUMENTED_LINE',
 'INSTRUMENTED_LOAD_SUPER_ATTR',
 'INSTRUMENTED_POP_JUMP_IF_FALSE',
 'INSTRUMENTED_POP_JUMP_IF_NONE',
 'INSTRUMENTED_POP_JUMP_IF_NOT_NONE',
 'INSTRUMENTED_POP_JUMP_IF_TRUE',
 'INSTRUMENTED_RESUME',
 'INSTRUMENTED_RETURN_CONST',
 'INSTRUMENTED_RETURN_VALUE',
 'INSTRUMENTED_YIELD_VALUE',
 'INTERPRETER_EXIT',
 'JUMP',
 'JUMP_NO_INTERRUPT',
 'LOAD_FROM_DICT_OR_DEREF',
 'LOAD_FROM_DICT_OR_GLOBALS',
 'LOAD_LOCALS',
 'LOAD_SUPER_METHOD',
 'LOAD_ZERO_SUPER_ATTR',
 'LOAD_ZERO_SUPER_METHOD',
 'POP_BLOCK',
 'RESERVED',
 'SETUP_ANNOTATIONS',
 'SETUP_CLEANUP',
 'SETUP_FINALLY',
 'SETUP_WITH',
 'STORE_FAST_MAYBE_NULL'}
{k for k, id in dis.opmap.items() if id < dis.opmap['INSTRUMENTED_LOAD_SUPER_ATTR']} - seen

gives

{'CACHE',
 'CALL_INTRINSIC_2',
 'CHECK_EG_MATCH',
 'CLEANUP_THROW',
 'INTERPRETER_EXIT',
 'LOAD_ASSERTION_ERROR',
 'LOAD_FROM_DICT_OR_DEREF',
 'LOAD_FROM_DICT_OR_GLOBALS',
 'LOAD_LOCALS',
 'RESERVED',
 'SETUP_ANNOTATIONS'}

so that (without CACHE and RESERVED) is much probably the list of things to think about for testing (LOAD_ASSERTION_ERROR) and implementations.

t-vi commented 1 month ago

CLEANUP_THROW is only seen in exception handling of yield_from, which would need improvments #815, LOAD_ASSERTION_ERROR is not seen in tests because PyTest rewriting asserts before execution of tests, INTERPRETER_EXIT is not generated directly but used for an interpreter trampoline (?) internally, CHECK_EG_MATCH is something we ignored for 3.11, too, we could eventually implement that (and look at whether we need anything for exception groups). LOAD_FROM_DICT_OR_DEREF, possibly LOAD_LOCALS and likely LOAD_FROM_DICT_OR_GLOBALS are used for class construction which we currently do not handle (because we call the opaque __build_class__ function instead of a lookaside that traces through the code). SETUP_ANNOTATIONS also seems about class (and module?) annotations. Here is a snippet where the dis shows a few of the opcodes in the class code object:

def fn():
    D = 1
    class A:
        B = 1
        C = 2 * B
        E:int = D
        def __init__(self, a: int):
            pass

    return A

CALL_INTRINSIC_2 and the remaining CALL_INTRINSIC_1 are about advanced typing. When we see this in the wild, we need to support it, but I didn't try running more than basic models yet.

To my mind, none of this is critical right now, I filed #815 about exceptions during YIELD.