Refactored Model Patcher Class #55

achew010 · 2024-07-18T11:01:51Z

Description

This PR addresses #44, it moves the ModelPatcher class to the Framework plugin so that all plugin patches can be managed and maintained under a common framework.

Items

Move all ModelPatcher functionality from fuse-ops-and-kernels to fms-acceleration
Replace local plugin patches with ModelPatcher functions
Unit tests to test ModelPatcher functionality

Benchmark Tests

We observe a general decrease in memory usage in this PR's benchmark results (memory plots are higher for reference), this led to some experiments running to completion when they ran out of memory previously.

fabianlim · 2024-07-18T13:04:17Z

plugins/fused-ops-and-kernels/src/fms_acceleration_foak/framework_plugin_fast_quantized_peft.py

-        model = patch_model(model, base_type=self._base_layer)
+        # wrapper function to register foak patches
+        from fms_acceleration_foak.models import load_foak_patches
+        load_foak_patches(base_type = self._base_layer)


is there a better name then load_foak_patches? You need to have more comments,

for what reason patching is required.

have a more descriptive name like registering_model_patches_for_blah_blah

dont import here, as there is no need to guard this import

plugins/framework/tests/test_model_patcher.py

plugins/framework/src/fms_acceleration/utils/test_utils.py

plugins/accelerated-peft/src/fms_acceleration_peft/autogptq_utils.py

plugins/framework/src/fms_acceleration/utils/test_utils.py

achew010 · 2024-07-21T11:38:49Z

@fabianlim This is the test plan that i'm using to write the unit test cases.

Test Plan

Objective:

Ensure ModelPatcher features produces the right outputs
Test the constraints of ModelPatcher, ModelPatcherRule, ModelPatcherTrigger and exceptions raised

1. ModelPatcherHistory

Test ModelPatcherHistory Construction
- test_mp_history_constructs_successfully()

2. ModelPatcherTrigger

Test Trigger Construction
- test_mp_trigger_constructs_successfully()
is_triggered
- Test Trigger checks respond as intended
  - test_mp_trigger_returns_correct_response()
combine_triggers
- test_correct_output_combine_mp_triggers()

3. ModelPatcherRule

Test ModelPatcherRule Construction
- test_mp_rule_constructs_successfully()

4. ModelPatcher

combine_functions
- test_correct_outputs_mp_combine_functions()
patch_target_module
- test_standalone_import_and_reload_function_replaces_indirect_module
register behaviour
- test_mp_registers_only_one_unique_rule()
patch scenarios
- Simple forward replacement
  - test_mp_rule_patches_forward
- import_and_reload replacement
  - test_MP_rule_import_and_reload_patches_downstream_module()
- forward builder replacement
  - test_MP_rule_patches_forward_with_builder_and_args
summary behaviour
- test_mp_history_converts_to_dataframe
load_patches behaviour
- test_registration_of_rules_through_package_imports

fabianlim

rename test_model_pactcher.py -> test_model_patcher_helpers.py
rename test_model_patcher2.py -> test_model_patcher.py
run an unassisted bench on 7b for related bech scenarios that use model patcher.

plugins/framework/src/fms_acceleration/model_patcher.py

…llback

… to allow triton access to global constexpr

fabianlim · 2024-07-29T04:31:46Z

plugins/framework/src/fms_acceleration/model_patcher.py

-        assert len(_with_reload) <= 1, "cannot have have at most one rule with reload"
+        # If there are multiple reload targets,
+        # ensure that their paths do not conflict as reloading same module might reset patches
+        if len(_with_reload)>1:


fabianlim · 2024-07-29T04:32:58Z

tox.ini

    packaging # this is required for flash-attn dep as fms_hf_tuning did not specify
    -e {toxinidir}/plugins/framework # install the framework here as the flash attention deps requires torch
 passenv = * # will pass the parent env, otherwise there are too many envs e.g. TRANSFORMERS that need to be set
+setenv = 
+    TRITON_ALLOW_NON_CONSTEXPR_GLOBALS=1


can you put a link to where the documentation says this must be set.

achew010 added 6 commits July 17, 2024 03:25

set main to track current plugin versions

7941ed7

move model_patcher to framework

4b871d0

replace local patching with model_patcher

3bf9a55

add additional unit tests

815b0c8

remove redundant patch function

7efbfed

shifted patch summary logging to framework plugin and patch id renames

33258ba

fabianlim reviewed Jul 18, 2024

View reviewed changes

plugins/framework/tests/test_model_patcher.py Outdated Show resolved Hide resolved

fabianlim reviewed Jul 19, 2024

View reviewed changes

plugins/framework/tests/test_model_patcher.py Outdated Show resolved Hide resolved

fabianlim reviewed Jul 19, 2024

View reviewed changes

plugins/framework/src/fms_acceleration/utils/test_utils.py Outdated Show resolved Hide resolved

fabianlim reviewed Jul 19, 2024

View reviewed changes

plugins/accelerated-peft/src/fms_acceleration_peft/autogptq_utils.py Show resolved Hide resolved

modified unit tests from PR comments

af7009c

fabianlim reviewed Jul 20, 2024

View reviewed changes

plugins/framework/src/fms_acceleration/utils/test_utils.py Show resolved Hide resolved

incremental refactor of unit tests

6b6fca9

achew010 changed the title ~~Refectored Model Patcher Class~~ Refactored Model Patcher Class Jul 22, 2024

achew010 added 3 commits July 23, 2024 02:20

changes to mp trigger unit tests

252a73c

additional changes to trigger unit tests

94e217e

adding MP Rule unit tests

a31bf6e

achew010 force-pushed the refactor/model-patcher branch from 785b971 to a31bf6e Compare July 23, 2024 11:48

achew010 and others added 3 commits July 24, 2024 04:37

add context manager to isolate patching unit tests

2683d9e

some fixes

748595c

clarified comments

9438aba

fabianlim force-pushed the refactor/model-patcher branch from 373e341 to 9438aba Compare July 25, 2024 03:45

achew010 marked this pull request as ready for review July 25, 2024 07:20

modelpatcher unit tests

8c825d9

achew010 force-pushed the refactor/model-patcher branch from 02e0998 to 8c825d9 Compare July 25, 2024 07:51

achew010 added 3 commits July 25, 2024 10:19

added forward_builder fn unit test

df95ece

lint changes

e653b80

more lint changes

e6f2284

fabianlim requested changes Jul 25, 2024

View reviewed changes

plugins/framework/src/fms_acceleration/model_patcher.py Outdated Show resolved Hide resolved

plugins/framework/src/fms_acceleration/model_patcher.py Show resolved Hide resolved

achew010 added 6 commits July 26, 2024 01:41

file renaming and added license headers on new files

736e706

added guard to patch model only if model exist in framework plugin ca…

7c302ba

…llback

replaced buggy partial wrapping on ModelPatcher.patch and set tox env…

cd253b3

… to allow triton access to global constexpr

additional linting

1d498e0

shifted patch trigger to main framework class

a4f8800

additional modifications to foak patch rules

ac31192

achew010 force-pushed the refactor/model-patcher branch from aff75fd to ac31192 Compare July 29, 2024 03:34

linting

8895cad

fabianlim reviewed Jul 29, 2024

View reviewed changes

achew010 force-pushed the refactor/model-patcher branch from c96c4c3 to 02a92e5 Compare July 29, 2024 05:02

additional changes from comments

f6848a7

achew010 force-pushed the refactor/model-patcher branch from 02a92e5 to f6848a7 Compare July 29, 2024 05:19

achew010 added 2 commits July 29, 2024 06:17

fixes to mp unit test

5e535b2

updated with new benchmark results

c204c86

fabianlim merged commit b6c1455 into foundation-model-stack:main Jul 29, 2024
4 checks passed

fabianlim mentioned this pull request Jul 30, 2024

Extract Out Model Patcher to Framework #44

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactored Model Patcher Class #55

Refactored Model Patcher Class #55

achew010 commented Jul 18, 2024 •

edited

Loading

fabianlim Jul 18, 2024 •

edited

Loading

fabianlim Jul 19, 2024

achew010 commented Jul 21, 2024

fabianlim left a comment •

edited

Loading

fabianlim Jul 29, 2024

fabianlim Jul 29, 2024

Refactored Model Patcher Class #55

Refactored Model Patcher Class #55

Conversation

achew010 commented Jul 18, 2024 • edited Loading

Description

Items

Benchmark Tests

fabianlim Jul 18, 2024 • edited Loading

Choose a reason for hiding this comment

fabianlim Jul 19, 2024

Choose a reason for hiding this comment

achew010 commented Jul 21, 2024

Test Plan

fabianlim left a comment • edited Loading

Choose a reason for hiding this comment

fabianlim Jul 29, 2024

Choose a reason for hiding this comment

fabianlim Jul 29, 2024

Choose a reason for hiding this comment

achew010 commented Jul 18, 2024 •

edited

Loading

fabianlim Jul 18, 2024 •

edited

Loading

fabianlim left a comment •

edited

Loading