File size: 1,652 Bytes
835cf28
 
 
 
 
 
 
 
 
 
554e4ad
835cf28
 
 
554e4ad
835cf28
 
 
 
 
 
 
 
554e4ad
 
835cf28
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import spaces
import torch
from typing import Callable, ParamSpec, Any
from spaces.zero.torch.aoti import ZeroGPUCompiledModel, ZeroGPUWeights

P = ParamSpec('P')

# Z-Image sequences vary based on resolution.
TRANSFORMER_SEQ_LENGTH = torch.export.Dim('seq_length', min=256, max=16384)

# forward(self, x, attn_mask, freqs_cis, adaln_input)
TRANSFORMER_DYNAMIC_SHAPES = (
    {1: TRANSFORMER_SEQ_LENGTH}, # x
    {1: TRANSFORMER_SEQ_LENGTH}, # attn_mask
    {1: TRANSFORMER_SEQ_LENGTH}, # freqs_cis
    None,                        # adaln_input
)

INDUCTOR_CONFIGS = {
    'conv_1x1_as_mm': True,
    'epilogue_fusion': False,
    'coordinate_descent_tuning': True,
    'coordinate_descent_check_all_directions': True,
    'max_autotune': True,
    'triton.cudagraphs': True,
}

def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):

    target_layers = pipeline.transformer.layers

    @spaces.GPU(duration=1500)
    def compile_block():
        block = target_layers[0]
        
        with spaces.aoti_capture(block) as call:
            pipeline(*args, **kwargs)

        dynamic_shapes = TRANSFORMER_DYNAMIC_SHAPES

        with torch.no_grad():
            exported = torch.export.export(
                mod=block,
                args=call.args,
                kwargs=call.kwargs,
                dynamic_shapes=dynamic_shapes,
            )

        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file

    archive_file = compile_block()

    for block in target_layers:
        weights = ZeroGPUWeights(block.state_dict())
        block.forward = ZeroGPUCompiledModel(archive_file, weights)