FilesExpand file tree

master

/

pinned_memory.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

128 lines (108 loc) · 4.57 KB

master

/

pinned_memory.py

Copy path

Top

File metadata and controls

• Code
• Blame

128 lines (108 loc) · 4.57 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

import bisect

import comfy.model_management

import comfy.memory_management

import comfy.utils

import comfy_aimdo.host_buffer

import comfy_aimdo.torch

import torch

from comfy.cli_args import args

def _add_to_bucket(module, module_pin, buckets, size, priority):

bucket = buckets.setdefault(size, [])

entry = [-priority, 0, module]

entry[1] = id(entry)

bisect.insort(bucket, entry)

module_pin["balancer_entry"] = entry

def _steal_pin(module, stack, buckets, size, priority, subset):

bucket = buckets.get(size)

if bucket is None:

return False

while bucket and bucket[-1][-1] is None:

bucket.pop()

if not bucket:

del buckets[size]

return False

if priority <= -bucket[-1][0]:

return False

*_, victim = bucket.pop()

module_pin = module._pins[subset]

victim_pin = victim._pins[subset]

module_pin["pin"] = victim_pin["pin"]

module_pin["registered"] = victim_pin["registered"]

module_pin["stack_index"] = victim_pin["stack_index"]

stack_index = module_pin["stack_index"]

stack[stack_index] = (module, stack[stack_index][1])

victim_pin["registered"] = False

del victim_pin["pin"]

del victim_pin["stack_index"]

del victim_pin["balancer_entry"]

_add_to_bucket(module, module_pin, buckets, size, priority)

return True

def get_pin(module, subset="weights"):

pins = module.__dict__.get("_pins")

module_pin = None if pins is None else pins.get(subset)

pin = None if module_pin is None else module_pin.get("pin")

if pin is None or module_pin["registered"] or args.disable_pinned_memory:

return pin

_, _, stack_split, pinned_size, *_ = module._pin_state[subset]

size = pin.nbytes

comfy.model_management.ensure_pin_registerable(size)

if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:

comfy.model_management.discard_cuda_async_error()

return pin

module_pin["registered"] = True

stack_split[0] = max(stack_split[0], module_pin["stack_index"])

comfy.model_management.TOTAL_PINNED_MEMORY += size

pinned_size[0] += size

return pin

def pin_memory(module, subset="weights", size=None):

pin_state = module._pin_state

if args.disable_pinned_memory:

return

pin = get_pin(module, subset)

if pin is not None:

return

pins = module.__dict__.setdefault("_pins", {})

module_pin = pins.setdefault(subset, {})

hostbuf, stack, stack_split, pinned_size, counter, buckets = pin_state[subset]

if size is None:

size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])

registerable_size = size

loaded = subset.endswith("-loaded")

fast = subset.endswith("-fast")

priority = module_pin.get("balancer_priority")

if priority is None:

priority = comfy.utils.bit_reverse_range(counter[0], 16)

counter[0] += 1

module_pin["balancer_priority"] = priority

comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)

if (not comfy.model_management.ensure_pin_budget(size, loaded=loaded) or

not comfy.model_management.ensure_pin_registerable(registerable_size, evict_active=not fast)):

return _steal_pin(module, stack, buckets, size, priority, subset)

offset = hostbuf.size

extended = False

try:

hostbuf.extend(size=size, register=False)

extended = True

pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)[offset:offset + size]

pin.untyped_storage()._comfy_hostbuf = hostbuf

if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:

comfy.model_management.discard_cuda_async_error()

comfy.model_management.free_registrations(size, evict_active=not fast)

if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:

comfy.model_management.discard_cuda_async_error()

del pin

hostbuf.truncate(offset, do_unregister=False)

return _steal_pin(module, stack, buckets, size, priority, subset)

except RuntimeError:

if extended:

hostbuf.truncate(offset, do_unregister=False)

return _steal_pin(module, stack, buckets, size, priority, subset)

module_pin["pin"] = pin

stack.append((module, offset))

module_pin["registered"] = True

module_pin["stack_index"] = len(stack) - 1

stack_split[0] = max(stack_split[0], module_pin["stack_index"])

comfy.model_management.TOTAL_PINNED_MEMORY += size

pinned_size[0] += size

_add_to_bucket(module, module_pin, buckets, size, priority)

return True