awslabs / awslabs/ratex

Native trace problem OPs

Open
#30 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
23
Forks
13
PR merge metrics
No merged PRs in 30d

Description

I finished tracing the model (forward & backward) w/ old package. Here are the ops that are currently giving me problems:

| OP | Problems | Potential causes | Tried/anticipated solution(s) | Resolved
|--------------|--------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------|-------|
| matmul | Current implementation causes tracing to be nullified | When changing dimensions of matrix, reshape/squeeze/unsqueeze/expand/view ops overwrite/remove tracing information from tensors. | Research & implement a bypass option for matmul or custom autograd implementation (experiment with test case) | This issue has been resolved|
| matmul | Current implementation does not support higher dimensional (4D+) | When reshaping higher dimensional matrix to lower, the values multiplied are inconsistent. | Tried tvm transpose/reshape algorithm. Other options include: implement/lower to einsum, employ a folding algorithm like aten, model off XLA | This issue has been resolved|
| embedding_dx | Midway gradient shape is invalid (at index 0 - got [512] but expected shape compatible with [512, 16]) | RAF embedding output updated weight tensor only as indices tensor is same. Need to update RAF to return full embedding table update or somehow merge indices tensor & output in Ratex | Currently have work around by unsqueeze & repeat OP to expand [512] -> [512,16]. However, we will need to implement proper solution and trace ltc graph/monitor gradient flow w/ test cases. For current progress & workaround implementation, check [EmbeddingDxBranch](https://github.com/awslabs/ratex/compare/main...anadeem2:ratex:native_trace?expand=1) | On going|
| layer_norm | Test case currently fails with mismatched values | Maybe raf_layer_norm & pytorch layer_norm have different naming. | Try running test case with lowering to raf_layer_norm_train_dx or other tvm offerings. It's possible that the mismatched values are normal, but we need to investigate/research to determine why the values are not identical. Maximum absolute diference is around ~1-2 while relative total difference is > 500 | On going|

From my analysis so far, it does not look like there are missing operators in the backward trace. The bigger challenge and the next step for me will be to identified decomposed ops & trace/lower them as complete ops (ex. index_put traced to decomposed embedding_dx).

I also ran BFloat16 model, the only trouble ops (both RAF/backend issues) are:
| OP | Problems | Potential causes | Tried/anticipated solution(s) | Resolved |
|------------|-------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------|-----|
| Dropout+dx | [F tensor.cpp:568] Check failed: handle->HasValue() Trying to access data while an async operation is in flight: f16[128,128] | Not sure, but dropout is supposed to return 3 outputs (tuple 3 (output, mask, reserve_space). Maybe somewhere bfloat16 is incompatible with one of these. | Tried RAF test case w/ float16. Error is resulting from RAF as raf.dropout does not support float16. | On going |
| Gelu+dx | gelu kernel is not implemented for float16/half. | Gelu approaches asymptote, so maybe float16 is not enough precision. | Tried RAF test case w/ float16. Failed all cases with errors being mismatched elements, and RuntimeError: "GeluKernelImpl" not implemented for 'Half' | On going |

I also tried to trace new package w/ layer_norm & native_batch_norm aggregated ops. I currently face graph fuser issues with that trace, and will return to this after full completion of old package.

There is also an issue with backward common.py test case. I particularly noticed while running matmul test case where the matmul input tensor grad would have cuda memory errors. I tried changing to default torch.randn(requires_grad=True) but got inflight on test case. Problem is in tensor initializations & memory for grad. Need to update common.py & rewrite backward verify step method. Not a problem in M5

I will continue to PR & implement OPs, however, if you have bandwidth or would like to contribute, please reply and take responsibility of ops. I can share my progress & work together as well.

Contributor guide

Open the contributing guide

Research direction

Review the unresolved operator table and start with the EmbeddingDxBranch linked in the issue, then inspect common.py for the backward verification and tensor-initialization problems. Run the relevant tracing and backward test cases for embedding_dx, layer_norm, dropout+dx, and gelu+dx. Done means the unresolved tracing, gradient-shape, value-mismatch, dtype, and verification failures have defined fixes and passing tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, pytorch
Domain
compilers, machine-learning, testing-qa
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.